{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "collapsed_sections": [
        "-N2TG4kkqIBx",
        "ZPQaYCG2Rzq_",
        "rTPWbL2BHRI4",
        "SiACUhFxDxOO",
        "XcdgkBphDucw",
        "BDKuuSYeuEnI",
        "bNuQvLoIOMuF",
        "EWbQyHxODjaH",
        "_Hkj0YUqD9y_",
        "R3LnR-SlEAhz",
        "IcBmQIHTECsg"
      ],
      "machine_shape": "hm",
      "gpuType": "A100",
      "toc_visible": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# **ATLAS** : Academic Task and Learning Agent System\n",
        "\n",
        "## **Overview**\n",
        "ATLAS demonstrates how to build an intelligent multi-agent system that transforms the way students manage their academic life. Using LangGraph's workflow framework, we'll create a network of specialized AI agents that work together to provide personalized academic support, from automated scheduling to intelligent lectures summarization.\n",
        "\n",
        "## **Motivation**\n",
        "Today's students face unprecedented challenges managing their academic workload alongside digital distractions and personal commitments. Traditional study planning tools often fall short because they:\n",
        "\n",
        "- Lack intelligent adaptation to individual learning styles\n",
        "- Don't integrate with students' existing digital ecosystems\n",
        "- Fail to provide context-aware assistance\n",
        "- Miss opportunities for proactive intervention\n",
        "\n",
        "**ATLAS** addresses these challenges through a sophisticated multi-agent architecture that combines advanced language models with structured workflows to deliver personalized academic support.\n",
        "##**Key Components**\n",
        "- Coordinator Agent: Orchestrates the interaction between specialized agents and manages the overall system state\n",
        "- Planner Agent: Handles calendar integration and schedule optimization\n",
        "- Notewriter Agent: Processes academic content and generates study materials\n",
        "- Advisor Agent: Provides personalized learning and time management advices\n",
        "\n",
        "## **Implementation Method**\n",
        "    ATLAS begins with a comprehensive initial assessment to understand each student's unique profile. The system conducts a thorough evaluation of learning preferences, cognitive styles, and current academic commitments while identifying specific challenges that require support. This information forms the foundation of a detailed student profile that drives personalized assistance throughout their academic journey.\n",
        "    At its core, ATLAS operates through a sophisticated multi-agent system architecture. The implementation leverages LangGraph's workflow framework to coordinate four specialized AI agents working in concert. The Coordinator Agent serves as the central orchestrator, managing workflow and ensuring seamless communication between components. The Planner Agent focuses on schedule optimization and time management, while the Notewriter Agent processes academic content and generates tailored study materials. The Advisor Agent rounds out the team by providing personalized guidance and support strategies.\n",
        "    The workflow orchestration implements a state management system that tracks student progress and coordinates agent activities. Using LangGraph's framework, the system maintains consistent communication channels between agents and defines clear transition rules for different academic scenarios. This structured approach ensures that each agent's specialized capabilities are deployed effectively to support student needs.\n",
        "    Learning process optimization forms a key part of the implementation. The system generates personalized study schedules that adapt to student preferences and energy patterns while creating customized learning materials that match individual learning styles. Real-time monitoring enables continuous adjustment of strategies based on student performance and engagement. The implementation incorporates proven learning techniques such as spaced repetition and active recall, automatically adjusting their application based on observed effectiveness.\n",
        "    Resource management and integration extend the system's capabilities through connections with external academic tools and platforms. ATLAS synchronizes with academic calendars, integrates with digital learning environments, and coordinates access to additional educational resources. This comprehensive integration ensures students have seamless access to all necessary tools and materials within their personalized academic support system.\n",
        "    The implementation maintains flexibility through continuous adaptation and improvement mechanisms. By monitoring performance metrics and gathering regular feedback, the system refines its recommendations and adjusts support strategies. This creates a dynamic learning environment that evolves with each student's changing needs and academic growth.\n",
        "    Emergency and support protocols are woven throughout the implementation to provide immediate assistance when needed. The system includes mechanisms for detecting academic stress, managing approaching deadlines, and providing intervention strategies during challenging periods. These protocols ensure students receive timely support while maintaining progress toward their academic goals.\n",
        "    Through this comprehensive implementation approach, ATLAS creates an intelligent, adaptive academic support system that grows increasingly effective at meeting each student's unique needs over time. The system's architecture enables seamless coordination between different support functions while maintaining focus on individual student success.\n",
        "\n",
        "## **Conclusion**\n",
        "ATLAS : Academic Task and Learning Agent System demonstrates the potential of combining language models with structured workflows to create an effective educational support system. By breaking down the academic support process into discrete steps and leveraging AI capabilities, we can provide personalized assistance that adapts to each student's needs. This approach opens up new possibilities for AI-assisted learning and academic success.\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ],
      "metadata": {
        "id": "zjqhTSGYpz_f"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Agents Design**"
      ],
      "metadata": {
        "id": "HOHVf968dCW-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "![agent_design.png]()![Screenshot 2024-11-18 110255.png]()"
      ],
      "metadata": {
        "id": "4wSHp-kYSNon"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Imports\n",
        "\n",
        "Install and import all necessary modules and libraries"
      ],
      "metadata": {
        "id": "-N2TG4kkqIBx"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%capture\n",
        "!pip install langgraph langchain langchain-openai openai python-dotenv"
      ],
      "metadata": {
        "id": "OUsHysWM1wne"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Graph Visualization\n",
        "%%capture\n",
        "!sudo apt-get install python3-dev graphviz libgraphviz-dev pkg-config\n",
        "!pip install graphviz\n",
        "!pip install pygraphviz"
      ],
      "metadata": {
        "collapsed": true,
        "id": "3Xy1v8wZaBQB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Utilities\n",
        "import operator\n",
        "from functools import reduce\n",
        "from typing import Annotated, List, Dict, TypedDict, Literal, Optional, Callable, Set, Tuple, Any, Union, TypeVar\n",
        "from datetime import datetime, timezone, timedelta\n",
        "import asyncio\n",
        "from pydantic import BaseModel, Field\n",
        "from operator import add\n",
        "from IPython.display import Image, display\n",
        "from google.colab import files\n",
        "import json\n",
        "import re\n",
        "import os\n",
        "# Core imports\n",
        "from openai import OpenAI, AsyncOpenAI\n",
        "from langchain_core.messages import HumanMessage, SystemMessage, BaseMessage\n",
        "from langchain.prompts import PromptTemplate\n",
        "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
        "from langgraph.graph import StateGraph, Graph, END, START\n",
        "\n",
        "\n",
        "# Pretty Markdown Output\n",
        "from rich.console import Console\n",
        "from rich.markdown import Markdown\n",
        "from rich.panel import Panel\n",
        "from rich.text import Text\n",
        "from rich import box\n",
        "from rich.style import Style"
      ],
      "metadata": {
        "id": "Z2onJd2yqhlN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "## API Configuration\n",
        "```\n",
        "Set up the API keys for the LLM provider (Nemotron-4-340B).\n",
        "\n",
        "For Google Colab:\n",
        "1. Add your API key to Colab secrets\n",
        "2. Name the secret 'NEMOTRON_4_340B_INSTRUCT_KEY'\n",
        "\n",
        "For local development:\n",
        "1. Create a .env file\n",
        "2. Add: NEMOTRON_4_340B_INSTRUCT_KEY=your_api_key\n",
        "```\n"
      ],
      "metadata": {
        "id": "ZPQaYCG2Rzq_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "NEMOTRON_4_340B_INSTRUCT_KEY = None #initialize global variable"
      ],
      "metadata": {
        "id": "JOocnKwUM-vg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def configure_api_keys():\n",
        "    \"\"\"Configure and verify API keys for LLM services.\"\"\"\n",
        "    # load_dotenv()\n",
        "    # api_key = os.getenv(\"NEMOTRON_4_340B_INSTRUCT_KEY\")\n",
        "\n",
        "    #Check Google Colab secrets\n",
        "    from google.colab import userdata\n",
        "    global NEMOTRON_4_340B_INSTRUCT_KEY\n",
        "    NEMOTRON_4_340B_INSTRUCT_KEY = userdata.get('NEMOTRON_4_340B_INSTRUCT_KEY')\n",
        "    # Set environment variable\n",
        "    os.environ['NEMOTRON_4_340B_INSTRUCT_KEY'] = NEMOTRON_4_340B_INSTRUCT_KEY\n",
        "    # Print configuration status\n",
        "    is_configured = bool(os.getenv(\"NEMOTRON_4_340B_INSTRUCT_KEY\"))\n",
        "    print(f\"API Key configured: {is_configured}\")\n",
        "    return is_configured\n",
        "\n",
        "api_configured = configure_api_keys()\n",
        "if not api_configured:\n",
        "    print(\"\\nAPI key not found. Please ensure you have:\")\n",
        "    print(\"1. Set up your API key in Google Colab secrets, or\")\n",
        "    print(\"2. Created a .env file with NEMOTRON_4_340B_INSTRUCT_KEY\")"
      ],
      "metadata": {
        "id": "I9p1Ng73qd_u",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "c8be9fd0-a8b8-407f-9fe1-94787398967d"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "API Key configured: True\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## State Definition\n",
        "\n",
        "Define the AcademicState class to hold the workflow's state."
      ],
      "metadata": {
        "id": "rTPWbL2BHRI4"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "T = TypeVar('T')\n",
        "\n",
        "def dict_reducer(dict1: Dict[str, Any], dict2: Dict[str, Any]) -> Dict[str, Any]:\n",
        "    \"\"\"\n",
        "    Merge two dictionaries recursively\n",
        "\n",
        "    Example:\n",
        "    dict1 = {\"a\": {\"x\": 1}, \"b\": 2}\n",
        "    dict2 = {\"a\": {\"y\": 2}, \"c\": 3}\n",
        "    result = {\"a\": {\"x\": 1, \"y\": 2}, \"b\": 2, \"c\": 3}\n",
        "    \"\"\"\n",
        "    merged = dict1.copy()\n",
        "    for key, value in dict2.items():\n",
        "        if key in merged and isinstance(merged[key], dict) and isinstance(value, dict):\n",
        "            merged[key] = dict_reducer(merged[key], value)\n",
        "        else:\n",
        "            merged[key] = value\n",
        "    return merged"
      ],
      "metadata": {
        "id": "vUcVfdq5NrWM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "class AcademicState(TypedDict):\n",
        "    \"\"\"Master state container for the academic assistance system\"\"\"\n",
        "    #  messages: Annotated[List[BaseMessage], add]   # Conversation history\n",
        "    #  profile: dict                                 # Student information\n",
        "    #  calendar: dict                                # Scheduled events\n",
        "    #  tasks: dict                                   # To-do items and assignments\n",
        "    #  results: Dict[str, Any]                       # Operation outputs\n",
        "    messages: Annotated[List[BaseMessage], add]   # Conversation history\n",
        "    profile: Annotated[Dict, dict_reducer]                 # Student information\n",
        "    calendar: Annotated[Dict, dict_reducer]                # Scheduled events\n",
        "    tasks: Annotated[Dict, dict_reducer]                   # To-do items and assignments\n",
        "    results: Annotated[Dict[str, Any], dict_reducer]       # Operation outputs"
      ],
      "metadata": {
        "id": "cdpWtKAPJwQP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# LLM Initialization\n"
      ],
      "metadata": {
        "id": "vnIAnM_iDzDH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Key Differences:**\n",
        "```\n",
        "1. Concurrency Model\n",
        "  - AsyncOpenAI: Asynchronous operations using `async/await`\n",
        "  - OpenAI: Synchronous operations that block execution\n",
        "\n",
        "2. Use Cases\n",
        "  - AsyncOpenAI: High throughput, non-blocking operations\n",
        "  - OpenAI: Simple sequential requests, easier debugging\n",
        "```"
      ],
      "metadata": {
        "id": "65TodIXnWLd-"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class LLMConfig:\n",
        "    \"\"\"Configuration settings for the LLM.\"\"\"\n",
        "    base_url: str = \"https://integrate.api.nvidia.com/v1\"\n",
        "    model: str = \"nvidia/nemotron-4-340b-instruct\"\n",
        "    max_tokens: int = 1024\n",
        "    default_temp: float = 0.5\n",
        "\n",
        "class NeMoLLaMa:\n",
        "  \"\"\"\n",
        "  A class to interact with NVIDIA's nemotron-4-340b-instruct model through their API\n",
        "  This implementation uses AsyncOpenAI client for asynchronous operations\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, api_key: str):\n",
        "    \"\"\"Initialize NeMoLLaMa with API key.\n",
        "\n",
        "    Args:\n",
        "        api_key (str): NVIDIA API authentication key\n",
        "    \"\"\"\n",
        "    self.config = LLMConfig()\n",
        "    self.client = AsyncOpenAI(\n",
        "        base_url=self.config.base_url,\n",
        "        api_key=api_key\n",
        "    )\n",
        "    self._is_authenticated = False\n",
        "\n",
        "  async def check_auth(self) -> bool:\n",
        "      \"\"\"Verify API authentication with test request.\n",
        "\n",
        "      Returns:\n",
        "          bool: Authentication status\n",
        "\n",
        "      Example:\n",
        "          >>> is_valid = await llm.check_auth()\n",
        "          >>> print(f\"Authenticated: {is_valid}\")\n",
        "      \"\"\"\n",
        "      test_message = [{\"role\": \"user\", \"content\": \"test\"}]\n",
        "      try:\n",
        "          await self.agenerate(test_message, temperature=0.1)\n",
        "          self._is_authenticated = True\n",
        "          return True\n",
        "      except Exception as e:\n",
        "          print(f\"❌ Authentication failed: {str(e)}\")\n",
        "          return False\n",
        "\n",
        "  async def agenerate(\n",
        "      self,\n",
        "      messages: List[Dict],\n",
        "      temperature: Optional[float] = None\n",
        "  ) -> str:\n",
        "      \"\"\"Generate text using NeMo LLaMa model.\n",
        "\n",
        "      Args:\n",
        "          messages: List of message dicts with 'role' and 'content'\n",
        "          temperature: Sampling temperature (0.0 to 1.0, default from config)\n",
        "\n",
        "      Returns:\n",
        "          str: Generated text response\n",
        "\n",
        "      Example:\n",
        "          >>> messages = [\n",
        "          ...     {\"role\": \"system\", \"content\": \"You are a helpful assistant\"},\n",
        "          ...     {\"role\": \"user\", \"content\": \"Plan my study schedule\"}\n",
        "          ... ]\n",
        "          >>> response = await llm.agenerate(messages, temperature=0.7)\n",
        "      \"\"\"\n",
        "      completion = await self.client.chat.completions.create(\n",
        "          model=self.config.model,\n",
        "          messages=messages,\n",
        "          temperature=temperature or self.config.default_temp,\n",
        "          max_tokens=self.config.max_tokens,\n",
        "          stream=False\n",
        "      )\n",
        "      return completion.choices[0].message.content"
      ],
      "metadata": {
        "id": "OxrqtnzI2fUk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### DataManager\n",
        "\n",
        "    A centralized data management system for AI agents to handle multiple data sources.\n",
        "    \n",
        "    This class serves as a unified interface for accessing and managing different types of\n",
        "    structured data (profiles, calendars, tasks) that an AI agent might need to process.\n",
        "    It handles data loading, parsing, and provides methods for intelligent filtering and retrieval.\n",
        "    "
      ],
      "metadata": {
        "id": "SiACUhFxDxOO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class DataManager:\n",
        "\n",
        "    def __init__(self):\n",
        "        \"\"\"\n",
        "        Initialize data storage containers.\n",
        "        All data sources start as None until explicitly loaded through load_data().\n",
        "        \"\"\"\n",
        "        self.profile_data = None\n",
        "        self.calendar_data = None\n",
        "        self.task_data = None\n",
        "\n",
        "    def load_data(self, profile_json: str, calendar_json: str, task_json: str):\n",
        "        \"\"\"\n",
        "        Load and parse multiple JSON data sources simultaneously.\n",
        "\n",
        "        Args:\n",
        "            profile_json (str): JSON string containing user profile information\n",
        "            calendar_json (str): JSON string containing calendar events\n",
        "            task_json (str): JSON string containing task/todo items\n",
        "\n",
        "        Note: This method expects valid JSON strings. Any parsing errors will propagate up.\n",
        "        \"\"\"\n",
        "        self.profile_data = json.loads(profile_json)\n",
        "        self.calendar_data = json.loads(calendar_json)\n",
        "        self.task_data = json.loads(task_json)\n",
        "\n",
        "    def get_student_profile(self, student_id: str) -> Dict:\n",
        "        \"\"\"\n",
        "        Retrieve a specific student's profile using their unique identifier.\n",
        "\n",
        "        Args:\n",
        "            student_id (str): Unique identifier for the student\n",
        "\n",
        "        Returns:\n",
        "            Dict: Student profile data if found, None otherwise\n",
        "\n",
        "        Implementation Note:\n",
        "            Uses generator expression with next() for efficient search through profiles,\n",
        "            avoiding full list iteration when possible.\n",
        "        \"\"\"\n",
        "        if self.profile_data:\n",
        "            return next((p for p in self.profile_data[\"profiles\"]\n",
        "                        if p[\"id\"] == student_id), None)\n",
        "        return None\n",
        "\n",
        "    def parse_datetime(self, dt_str: str) -> datetime:\n",
        "        \"\"\"\n",
        "        Smart datetime parser that handles multiple formats and ensures UTC timezone.\n",
        "\n",
        "        Args:\n",
        "            dt_str (str): DateTime string in ISO format, with or without timezone\n",
        "\n",
        "        Returns:\n",
        "            datetime: Parsed datetime object in UTC timezone\n",
        "\n",
        "        Implementation Note:\n",
        "            Handles both timezone-aware and naive datetime strings by:\n",
        "            1. First attempting to parse with timezone information\n",
        "            2. Falling back to assuming UTC if no timezone is specified\n",
        "        \"\"\"\n",
        "        try:\n",
        "            # First attempt: Parse ISO format with timezone\n",
        "            dt = datetime.fromisoformat(dt_str.replace('Z', '+00:00'))\n",
        "            return dt.astimezone(timezone.utc)\n",
        "        except ValueError:\n",
        "            # Fallback: Assume UTC if no timezone provided\n",
        "            dt = datetime.fromisoformat(dt_str)\n",
        "            return dt.replace(tzinfo=timezone.utc)\n",
        "\n",
        "    def get_upcoming_events(self, days: int = 7) -> List[Dict]:\n",
        "        \"\"\"\n",
        "        Intelligently filter and retrieve upcoming calendar events within a specified timeframe.\n",
        "\n",
        "        Args:\n",
        "            days (int): Number of days to look ahead (default: 7)\n",
        "\n",
        "        Returns:\n",
        "            List[Dict]: List of upcoming events, chronologically ordered\n",
        "\n",
        "        Implementation Note:\n",
        "            - Uses UTC timestamps for consistent timezone handling\n",
        "            - Implements error handling for malformed event data\n",
        "            - Only includes events that start in the future up to the specified timeframe\n",
        "        \"\"\"\n",
        "        if not self.calendar_data:\n",
        "            return []\n",
        "\n",
        "        now = datetime.now(timezone.utc)\n",
        "        future = now + timedelta(days=days)\n",
        "\n",
        "        events = []\n",
        "        for event in self.calendar_data.get(\"events\", []):\n",
        "            try:\n",
        "                start_time = self.parse_datetime(event[\"start\"][\"dateTime\"])\n",
        "\n",
        "                if now <= start_time <= future:\n",
        "                    events.append(event)\n",
        "            except (KeyError, ValueError) as e:\n",
        "                print(f\"Warning: Could not process event due to {str(e)}\")\n",
        "                continue\n",
        "\n",
        "        return events\n",
        "\n",
        "    def get_active_tasks(self) -> List[Dict]:\n",
        "        \"\"\"\n",
        "        Retrieve and filter active tasks, enriching them with parsed datetime information.\n",
        "\n",
        "        Returns:\n",
        "            List[Dict]: List of active tasks with parsed due dates\n",
        "\n",
        "        Implementation Note:\n",
        "            - Filters for tasks that are:\n",
        "              1. Not completed (\"needsAction\" status)\n",
        "              2. Due in the future\n",
        "            - Enriches task objects with parsed datetime for easier processing\n",
        "            - Implements robust error handling for malformed task data\n",
        "        \"\"\"\n",
        "        if not self.task_data:\n",
        "            return []\n",
        "\n",
        "        now = datetime.now(timezone.utc)\n",
        "        active_tasks = []\n",
        "\n",
        "        for task in self.task_data.get(\"tasks\", []):\n",
        "            try:\n",
        "                due_date = self.parse_datetime(task[\"due\"])\n",
        "                if task[\"status\"] == \"needsAction\" and due_date > now:\n",
        "                    # Enrich task object with parsed datetime\n",
        "                    task[\"due_datetime\"] = due_date\n",
        "                    active_tasks.append(task)\n",
        "            except (KeyError, ValueError) as e:\n",
        "                print(f\"Warning: Could not process task due to {str(e)}\")\n",
        "                continue\n",
        "\n",
        "        return active_tasks"
      ],
      "metadata": {
        "id": "ECQKdxnl2owa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# usage\n",
        "llm = NeMoLLaMa(NEMOTRON_4_340B_INSTRUCT_KEY)\n",
        "data_manager = DataManager()\n",
        "print(llm)"
      ],
      "metadata": {
        "id": "gaKUz47LEOQY",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "e5066d6b-1baf-4c8d-e437-1a09857c6d0b"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "<__main__.NeMoLLaMa object at 0x7e3cc6158790>\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Agent Executor"
      ],
      "metadata": {
        "id": "78pTnTlSxF9i"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "    \n",
        "    Orchestrates the concurrent execution of multiple specialized AI agents.\n",
        "    \n",
        "    This class implements a sophisticated execution pattern that allows multiple AI agents\n",
        "    to work together, either sequentially or concurrently, based on a coordination analysis.\n",
        "    It handles agent initialization, concurrent execution, error handling, and fallback strategies.\n",
        "    "
      ],
      "metadata": {
        "id": "0sR3eTuUw7KA"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class AgentExecutor:\n",
        "\n",
        "\n",
        "    def __init__(self, llm):\n",
        "        \"\"\"\n",
        "        Initialize the executor with a language model and create agent instances.\n",
        "\n",
        "        Args:\n",
        "            llm: Language model instance to be used by all agents\n",
        "\n",
        "        Implementation Note:\n",
        "            - Creates a dictionary of specialized agents, each initialized with the same LLM\n",
        "            - Supports multiple agent types: PLANNER (default), NOTEWRITER, and ADVISOR\n",
        "            - Agents are instantiated once and reused across executions\n",
        "        \"\"\"\n",
        "        self.llm = llm\n",
        "        self.agents = {\n",
        "            \"PLANNER\": PlannerAgent(llm),      # Strategic planning agent\n",
        "            \"NOTEWRITER\": NoteWriterAgent(llm), # Documentation agent\n",
        "            \"ADVISOR\": AdvisorAgent(llm)        # Academic advice agent\n",
        "        }\n",
        "\n",
        "    async def execute(self, state: AcademicState) -> Dict:\n",
        "        \"\"\"\n",
        "        Orchestrates concurrent execution of multiple AI agents based on analysis results.\n",
        "\n",
        "        This method implements a sophisticated execution pattern:\n",
        "        1. Reads coordination analysis to determine required agents\n",
        "        2. Groups agents for concurrent execution\n",
        "        3. Executes agent groups in parallel\n",
        "        4. Handles failures gracefully with fallback mechanisms\n",
        "\n",
        "        Args:\n",
        "            state (AcademicState): Current academic state containing analysis results\n",
        "\n",
        "        Returns:\n",
        "            Dict: Consolidated results from all executed agents\n",
        "\n",
        "        Implementation Details:\n",
        "        ---------------------\n",
        "        1. Analysis Interpretation:\n",
        "           - Extracts coordination analysis from state\n",
        "           - Determines required agents and their concurrent execution groups\n",
        "\n",
        "        2. Concurrent Execution Pattern:\n",
        "           - Processes agents in groups that can run in parallel\n",
        "           - Uses asyncio.gather() for concurrent execution within each group\n",
        "           - Only executes agents that are both required and available\n",
        "\n",
        "        3. Result Management:\n",
        "           - Collects and processes results from each concurrent group\n",
        "           - Filters out failed executions (exceptions)\n",
        "           - Formats successful results into a structured output\n",
        "\n",
        "        4. Fallback Mechanisms:\n",
        "           - If no results are gathered, falls back to PLANNER agent\n",
        "           - Provides emergency fallback plan in case of complete failure\n",
        "\n",
        "        Error Handling:\n",
        "        --------------\n",
        "        - Catches and handles exceptions at multiple levels:\n",
        "          * Individual agent execution failures don't affect other agents\n",
        "          * System-level failures trigger emergency fallback\n",
        "        - Maintains system stability through graceful degradation\n",
        "        \"\"\"\n",
        "        try:\n",
        "            # Extract coordination analysis from state\n",
        "            analysis = state[\"results\"].get(\"coordinator_analysis\", {})\n",
        "\n",
        "            # Determine execution requirements\n",
        "            required_agents = analysis.get(\"required_agents\", [\"PLANNER\"])  # PLANNER as default\n",
        "            concurrent_groups = analysis.get(\"concurrent_groups\", [])       # Groups for parallel execution\n",
        "\n",
        "            # Initialize results container\n",
        "            results = {}\n",
        "\n",
        "            # Process each concurrent group sequentially\n",
        "            for group in concurrent_groups:\n",
        "                # Prepare concurrent tasks for current group\n",
        "                tasks = []\n",
        "                for agent_name in group:\n",
        "                    # Validate agent availability and requirement\n",
        "                    if agent_name in required_agents and agent_name in self.agents:\n",
        "                        tasks.append(self.agents[agent_name](state))\n",
        "\n",
        "                # Execute group tasks concurrently if any exist\n",
        "                if tasks:\n",
        "                    # Gather results from concurrent execution\n",
        "                    group_results = await asyncio.gather(*tasks, return_exceptions=True)\n",
        "\n",
        "                    # Process successful results only\n",
        "                    for agent_name, result in zip(group, group_results):\n",
        "                        if not isinstance(result, Exception):\n",
        "                            results[agent_name.lower()] = result\n",
        "\n",
        "            # Implement fallback strategy if no results were obtained\n",
        "            if not results and \"PLANNER\" in self.agents:\n",
        "                planner_result = await self.agents[\"PLANNER\"](state)\n",
        "                results[\"planner\"] = planner_result\n",
        "\n",
        "            print(\"agent_outputs\", results)\n",
        "\n",
        "            # Return structured results\n",
        "            return {\n",
        "                \"results\": {\n",
        "                    \"agent_outputs\": results\n",
        "                }\n",
        "            }\n",
        "\n",
        "        except Exception as e:\n",
        "            print(f\"Execution error: {e}\")\n",
        "            # Emergency fallback with minimal response\n",
        "            return {\n",
        "                \"results\": {\n",
        "                    \"agent_outputs\": {\n",
        "                        \"planner\": {\n",
        "                            \"plan\": \"Emergency fallback plan: Please try again or contact support.\"\n",
        "                        }\n",
        "                    }\n",
        "                }\n",
        "            }"
      ],
      "metadata": {
        "id": "sRkQgRaIwSdB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Agent Action and Output Models\n",
        "- Defines the structure for agent actions and outputs using Pydantic models.\n",
        "These models ensure type safety and validation for agent operations."
      ],
      "metadata": {
        "id": "0R4W2N9iNpcW"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class AgentAction(BaseModel):\n",
        "   \"\"\"\n",
        "   Model representing an agent's action decision.\n",
        "\n",
        "   Attributes:\n",
        "       action (str): The specific action to be taken (e.g., \"search_calendar\", \"analyze_tasks\")\n",
        "       thought (str): The reasoning process behind the action choice\n",
        "       tool (Optional[str]): The specific tool to be used for the action (if needed)\n",
        "       action_input (Optional[Dict]): Input parameters for the action\n",
        "\n",
        "   Example:\n",
        "       >>> action = AgentAction(\n",
        "       ...     action=\"search_calendar\",\n",
        "       ...     thought=\"Need to check schedule conflicts\",\n",
        "       ...     tool=\"calendar_search\",\n",
        "       ...     action_input={\"date_range\": \"next_week\"}\n",
        "       ... )\n",
        "   \"\"\"\n",
        "   action: str      # Required action to be performed\n",
        "   thought: str     # Reasoning behind the action\n",
        "   tool: Optional[str] = None        # Optional tool specification\n",
        "   action_input: Optional[Dict] = None  # Optional input parameters\n",
        "\n",
        "class AgentOutput(BaseModel):\n",
        "   \"\"\"\n",
        "   Model representing the output from an agent's action.\n",
        "\n",
        "   Attributes:\n",
        "       observation (str): The result or observation from executing the action\n",
        "       output (Dict): Structured output data from the action\n",
        "\n",
        "   Example:\n",
        "       >>> output = AgentOutput(\n",
        "       ...     observation=\"Found 3 free time slots next week\",\n",
        "       ...     output={\n",
        "       ...         \"free_slots\": [\"Mon 2PM\", \"Wed 10AM\", \"Fri 3PM\"],\n",
        "       ...         \"conflicts\": []\n",
        "       ...     }\n",
        "       ... )\n",
        "   \"\"\"\n",
        "   observation: str  # Result or observation from the action\n",
        "   output: Dict     # Structured output data"
      ],
      "metadata": {
        "id": "fyz_bExHuD-J"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# ReACT agent"
      ],
      "metadata": {
        "id": "BDKuuSYeuEnI"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "**What's actually is ReACT?**\n",
        "\n",
        "ReACT (Reasoning and Acting) is a framework that combines reasoning and acting in an iterative process.\n",
        "It enables LLMs to approach complex tasks by breaking them down into:\n",
        "\n",
        "1. **(Re)act**: Take an action based on observations and tools\n",
        "2. **(Re)ason**: Think about what to do next\n",
        "3. **(Re)flect**: Learn from the outcome\n",
        "\n",
        "Example Flow:\n",
        "- Thought: Need to check student's schedule for study time\n",
        "- Action: search_calendar\n",
        "- Observation: Found 2 free hours tomorrow morning\n",
        "- Thought: Student prefers morning study, this is optimal\n",
        "- Action: analyze_tasks\n",
        "- Observation: Has 3 pending assignments\n",
        "- Plan: Schedule morning study session for highest priority task"
      ],
      "metadata": {
        "id": "68hB7fJONtZO"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ReActAgent:\n",
        "  \"\"\"\n",
        "    Base class for ReACT-based agents implementing reasoning and action capabilities.\n",
        "\n",
        "    Features:\n",
        "    - Tool management for specific actions\n",
        "    - Few-shot learning examples\n",
        "    - Structured thought process\n",
        "    - Action execution framework\n",
        "  \"\"\"\n",
        "\n",
        "  def __init__(self, llm):\n",
        "      \"\"\"\n",
        "      Initialize the ReActAgent with language model and available tools\n",
        "\n",
        "      Args:\n",
        "          llm: Language model instance for agent operations\n",
        "      \"\"\"\n",
        "      self.llm = llm\n",
        "      # Storage for few-shot examples to guide the agent\n",
        "      self.few_shot_examples = []\n",
        "\n",
        "      # Dictionary of available tools with their corresponding methods\n",
        "      self.tools = {\n",
        "          \"search_calendar\": self.search_calendar,      # Calendar search functionality\n",
        "          \"analyze_tasks\": self.analyze_tasks,          # Task analysis functionality\n",
        "          \"check_learning_style\": self.check_learning_style,  # Learning style assessment\n",
        "          \"check_performance\": self.check_performance   # Academic performance checking\n",
        "      }\n",
        "\n",
        "  async def search_calendar(self, state: AcademicState) -> List[Dict]:\n",
        "      \"\"\"\n",
        "      Search for upcoming calendar events\n",
        "\n",
        "      Args:\n",
        "          state (AcademicState): Current academic state\n",
        "\n",
        "      Returns:\n",
        "          List[Dict]: List of upcoming calendar events\n",
        "      \"\"\"\n",
        "      # Get events from calendar or empty list if none exist\n",
        "      events = state[\"calendar\"].get(\"events\", [])\n",
        "      # Get current time in UTC\n",
        "      now = datetime.now(timezone.utc)\n",
        "      # Filter and return only future events\n",
        "      return [e for e in events if datetime.fromisoformat(e[\"start\"][\"dateTime\"]) > now]\n",
        "\n",
        "  async def analyze_tasks(self, state: AcademicState) -> List[Dict]:\n",
        "      \"\"\"\n",
        "      Analyze academic tasks from the current state\n",
        "\n",
        "      Args:\n",
        "          state (AcademicState): Current academic state\n",
        "\n",
        "      Returns:\n",
        "          List[Dict]: List of academic tasks\n",
        "      \"\"\"\n",
        "      # Return tasks or empty list if none exist\n",
        "      return state[\"tasks\"].get(\"tasks\", [])\n",
        "\n",
        "  async def check_learning_style(self, state: AcademicState) -> AcademicState:\n",
        "        \"\"\"\n",
        "        Retrieve student's learning style and study patterns\n",
        "\n",
        "        Args:\n",
        "            state (AcademicState): Current academic state\n",
        "\n",
        "        Returns:\n",
        "            AcademicState: Updated state with learning style analysis\n",
        "        \"\"\"\n",
        "        # Get user profile from state\n",
        "        profile = state[\"profile\"]\n",
        "\n",
        "        # Get learning preferences\n",
        "        learning_data = {\n",
        "            \"style\": profile.get(\"learning_preferences\", {}).get(\"learning_style\", {}),\n",
        "            \"patterns\": profile.get(\"learning_preferences\", {}).get(\"study_patterns\", {})\n",
        "        }\n",
        "\n",
        "        # Add to results in state\n",
        "        if \"results\" not in state:\n",
        "            state[\"results\"] = {}\n",
        "        state[\"results\"][\"learning_analysis\"] = learning_data\n",
        "\n",
        "        return state\n",
        "\n",
        "  async def check_performance(self, state: AcademicState) -> AcademicState:\n",
        "        \"\"\"\n",
        "        Check current academic performance across courses\n",
        "\n",
        "        Args:\n",
        "            state (AcademicState): Current academic state\n",
        "\n",
        "        Returns:\n",
        "            AcademicState: Updated state with performance analysis\n",
        "        \"\"\"\n",
        "        # Get user profile from state\n",
        "        profile = state[\"profile\"]\n",
        "\n",
        "        # Get course information\n",
        "        courses = profile.get(\"academic_info\", {}).get(\"current_courses\", [])\n",
        "\n",
        "        # Add to results in state\n",
        "        if \"results\" not in state:\n",
        "            state[\"results\"] = {}\n",
        "        state[\"results\"][\"performance_analysis\"] = {\"courses\": courses}\n",
        "\n",
        "        return state"
      ],
      "metadata": {
        "id": "hGwZJRv6uRmi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Coordinator Agent"
      ],
      "metadata": {
        "id": "1jVc4HeT7HU3"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "async def analyze_context(state: AcademicState) -> Dict:\n",
        "    \"\"\"\n",
        "    Analyzes the academic state context to inform coordinator decision-making.\n",
        "\n",
        "    This function performs comprehensive context analysis by:\n",
        "    1. Extracting student profile information\n",
        "    2. Analyzing calendar and task loads\n",
        "    3. Identifying relevant course context from the latest message\n",
        "    4. Gathering learning preferences and study patterns\n",
        "\n",
        "    Args:\n",
        "        state (AcademicState): Current academic state including profile, calendar, and tasks\n",
        "\n",
        "    Returns:\n",
        "        Dict: Structured analysis of the student's context for decision making\n",
        "\n",
        "    Implementation Notes:\n",
        "    ------------------\n",
        "    - Extracts information hierarchically using nested get() operations for safety\n",
        "    - Identifies current course context from the latest message content\n",
        "    - Provides default values for missing information to ensure stability\n",
        "    \"\"\"\n",
        "    # Extract main data components with safe navigation\n",
        "    profile = state.get(\"profile\", {})\n",
        "    calendar = state.get(\"calendar\", {})\n",
        "    tasks = state.get(\"tasks\", {})\n",
        "\n",
        "    # Extract course information and match with current request\n",
        "    courses = profile.get(\"academic_info\", {}).get(\"current_courses\", [])\n",
        "    current_course = None\n",
        "    request = state[\"messages\"][-1].content.lower()  # Latest message for context\n",
        "\n",
        "    # Identify relevant course from request content\n",
        "    for course in courses:\n",
        "        if course[\"name\"].lower() in request:\n",
        "            current_course = course\n",
        "            break\n",
        "\n",
        "    # Construct comprehensive context analysis\n",
        "    return {\n",
        "        \"student\": {\n",
        "            \"major\": profile.get(\"personal_info\", {}).get(\"major\", \"Unknown\"),\n",
        "            \"year\": profile.get(\"personal_info\", {}).get(\"academic_year\"),\n",
        "            \"learning_style\": profile.get(\"learning_preferences\", {}).get(\"learning_style\", {}),\n",
        "        },\n",
        "        \"course\": current_course,\n",
        "        \"upcoming_events\": len(calendar.get(\"events\", [])),  # Calendar load indicator\n",
        "        \"active_tasks\": len(tasks.get(\"tasks\", [])),        # Task load indicator\n",
        "        \"study_patterns\": profile.get(\"learning_preferences\", {}).get(\"study_patterns\", {})\n",
        "    }\n",
        "\n",
        "def parse_coordinator_response(response: str) -> Dict:\n",
        "    \"\"\"\n",
        "    Parses LLM coordinator response into structured analysis for agent execution.\n",
        "\n",
        "    This function implements a robust parsing strategy:\n",
        "    1. Starts with safe default configuration\n",
        "    2. Analyzes ReACT patterns in the response\n",
        "    3. Adjusts agent requirements and priorities based on content\n",
        "    4. Organizes concurrent execution groups\n",
        "\n",
        "    Args:\n",
        "        response (str): Raw LLM response text\n",
        "\n",
        "    Returns:\n",
        "        Dict: Structured analysis containing:\n",
        "            - required_agents: List of agents needed\n",
        "            - priority: Priority levels for each agent\n",
        "            - concurrent_groups: Groups of agents that can run together\n",
        "            - reasoning: Extracted reasoning for decisions\n",
        "\n",
        "    Implementation Notes:\n",
        "    ------------------\n",
        "    1. Default Configuration:\n",
        "       - Always includes PLANNER agent as baseline\n",
        "       - Sets basic priority and concurrent execution structure\n",
        "\n",
        "    2. Response Analysis:\n",
        "       - Looks for ReACT patterns (Thought/Decision structure)\n",
        "       - Identifies agent requirements from content keywords\n",
        "       - Extracts reasoning from thought section\n",
        "\n",
        "    3. Agent Configuration:\n",
        "       - NOTEWRITER triggered by note-taking related content\n",
        "       - ADVISOR triggered by guidance/advice related content\n",
        "       - Organizes concurrent execution groups based on dependencies\n",
        "\n",
        "    4. Error Handling:\n",
        "       - Provides fallback configuration if parsing fails\n",
        "       - Maintains system stability through default values\n",
        "    \"\"\"\n",
        "    try:\n",
        "        # Initialize with safe default configuration\n",
        "        analysis = {\n",
        "            \"required_agents\": [\"PLANNER\"],         # PLANNER is always required\n",
        "            \"priority\": {\"PLANNER\": 1},             # Base priority structure\n",
        "            \"concurrent_groups\": [[\"PLANNER\"]],     # Default execution group\n",
        "            \"reasoning\": \"Default coordination\"      # Default reasoning\n",
        "        }\n",
        "\n",
        "        # Parse ReACT patterns for advanced coordination\n",
        "        if \"Thought:\" in response and \"Decision:\" in response:\n",
        "            # Check for NOTEWRITER requirements\n",
        "            if \"NoteWriter\" in response or \"note\" in response.lower():\n",
        "                analysis[\"required_agents\"].append(\"NOTEWRITER\")\n",
        "                analysis[\"priority\"][\"NOTEWRITER\"] = 2\n",
        "                # NOTEWRITER can run concurrently with PLANNER\n",
        "                analysis[\"concurrent_groups\"] = [[\"PLANNER\", \"NOTEWRITER\"]]\n",
        "\n",
        "            # Check for ADVISOR requirements\n",
        "            if \"Advisor\" in response or \"guidance\" in response.lower():\n",
        "                analysis[\"required_agents\"].append(\"ADVISOR\")\n",
        "                analysis[\"priority\"][\"ADVISOR\"] = 3\n",
        "                # ADVISOR typically runs after initial planning\n",
        "\n",
        "            # Extract and store reasoning from thought section\n",
        "            thought_section = response.split(\"Thought:\")[1].split(\"Action:\")[0].strip()\n",
        "            analysis[\"reasoning\"] = thought_section\n",
        "\n",
        "        return analysis\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"Parse error: {str(e)}\")\n",
        "        # Fallback to safe default configuration\n",
        "        return {\n",
        "            \"required_agents\": [\"PLANNER\"],\n",
        "            \"priority\": {\"PLANNER\": 1},\n",
        "            \"concurrent_groups\": [[\"PLANNER\"]],\n",
        "            \"reasoning\": \"Fallback due to parse error\"\n",
        "        }"
      ],
      "metadata": {
        "id": "6C3L-DSVvh3H"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Define Coordinator Agent Prompt with ReACT Prompting"
      ],
      "metadata": {
        "id": "ZfLKDqk8zo0A"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "COORDINATOR_PROMPT =\"\"\"You are a Coordinator Agent using ReACT framework to orchestrate multiple academic support agents.\n",
        "\n",
        "        AVAILABLE AGENTS:\n",
        "        • PLANNER: Handles scheduling and time management\n",
        "        • NOTEWRITER: Creates study materials and content summaries\n",
        "        • ADVISOR: Provides personalized academic guidance\n",
        "\n",
        "        PARALLEL EXECUTION RULES:\n",
        "        1. Group compatible agents that can run concurrently\n",
        "        2. Maintain dependencies between agent executions\n",
        "        3. Coordinate results from parallel executions\n",
        "\n",
        "        REACT PATTERN:\n",
        "        Thought: [Analyze request complexity and required support types]\n",
        "        Action: [Select optimal agent combination]\n",
        "        Observation: [Evaluate selected agents' capabilities]\n",
        "        Decision: [Finalize agent deployment plan]\n",
        "\n",
        "        ANALYSIS POINTS:\n",
        "        1. Task Complexity and Scope\n",
        "        2. Time Constraints\n",
        "        3. Resource Requirements\n",
        "        4. Learning Style Alignment\n",
        "        5. Support Type Needed\n",
        "\n",
        "        CONTEXT:\n",
        "        Request: {request}\n",
        "        Student Context: {context}\n",
        "\n",
        "        FORMAT RESPONSE AS:\n",
        "        Thought: [Analysis of academic needs and context]\n",
        "        Action: [Agent selection and grouping strategy]\n",
        "        Observation: [Expected workflow and dependencies]\n",
        "        Decision: [Final agent deployment plan with rationale]\n",
        "        \"\"\""
      ],
      "metadata": {
        "id": "ajOjzQCGz6kd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "async def coordinator_agent(state: AcademicState) -> Dict:\n",
        "    \"\"\"\n",
        "    Primary coordinator agent that orchestrates multiple academic support agents using ReACT framework.\n",
        "\n",
        "    This agent implements a sophisticated coordination strategy:\n",
        "    1. Analyzes academic context and student needs\n",
        "    2. Uses ReACT framework for structured decision making\n",
        "    3. Coordinates parallel agent execution\n",
        "    4. Handles fallback scenarios\n",
        "\n",
        "    Args:\n",
        "        state (AcademicState): Current academic state including messages and context\n",
        "\n",
        "    Returns:\n",
        "        Dict: Coordination analysis including required agents, priorities, and execution groups\n",
        "\n",
        "    Implementation Notes:\n",
        "    -------------------\n",
        "    1. ReACT Framework Implementation:\n",
        "       - Thought: Analysis phase\n",
        "       - Action: Agent selection phase\n",
        "       - Observation: Capability evaluation\n",
        "       - Decision: Final execution planning\n",
        "\n",
        "    2. Agent Coordination Strategy:\n",
        "       - Manages three specialized agents:\n",
        "         * PLANNER: Core scheduling agent\n",
        "         * NOTEWRITER: Content creation agent\n",
        "         * ADVISOR: Academic guidance agent\n",
        "\n",
        "    3. Parallel Execution Management:\n",
        "       - Groups compatible agents\n",
        "       - Maintains execution dependencies\n",
        "       - Coordinates parallel workflows\n",
        "    \"\"\"\n",
        "    try:\n",
        "        # Analyze current context and extract latest query\n",
        "        context = await analyze_context(state)\n",
        "        query = state[\"messages\"][-1].content\n",
        "\n",
        "        # Define the ReACT-based coordination prompt\n",
        "        prompt = COORDINATOR_PROMPT\n",
        "\n",
        "        # Generate coordination plan using LLM\n",
        "        response = await llm.agenerate([\n",
        "            {\"role\": \"system\", \"content\": prompt.format(\n",
        "                request=query,\n",
        "                context=json.dumps(context, indent=2)\n",
        "            )}\n",
        "        ])\n",
        "\n",
        "        # Parse response and structure coordination analysis\n",
        "        analysis = parse_coordinator_response(response)\n",
        "        return {\n",
        "            \"results\": {\n",
        "                \"coordinator_analysis\": {\n",
        "                    \"required_agents\": analysis.get(\"required_agents\", [\"PLANNER\"]),\n",
        "                    \"priority\": analysis.get(\"priority\", {\"PLANNER\": 1}),\n",
        "                    \"concurrent_groups\": analysis.get(\"concurrent_groups\", [[\"PLANNER\"]]),\n",
        "                    \"reasoning\": response\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"Coordinator error: {e}\")\n",
        "        # Fallback to basic planning configuration\n",
        "        return {\n",
        "            \"results\": {\n",
        "                \"coordinator_analysis\": {\n",
        "                    \"required_agents\": [\"PLANNER\"],\n",
        "                    \"priority\": {\"PLANNER\": 1},\n",
        "                    \"concurrent_groups\": [[\"PLANNER\"]],\n",
        "                    \"reasoning\": \"Error in coordination. Falling back to planner.\"\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "\n",
        "def parse_coordinator_response(response: str) -> Dict:\n",
        "    \"\"\"\n",
        "    Parses LLM response into structured coordination analysis.\n",
        "\n",
        "    This function:\n",
        "    1. Starts with default safe configuration\n",
        "    2. Analyzes ReACT pattern responses\n",
        "    3. Identifies required agents and priorities\n",
        "    4. Structures concurrent execution groups\n",
        "\n",
        "    Args:\n",
        "        response (str): Raw LLM response following ReACT pattern\n",
        "\n",
        "    Returns:\n",
        "        Dict: Structured analysis for agent execution\n",
        "\n",
        "    Implementation Notes:\n",
        "    -------------------\n",
        "    1. Default Configuration:\n",
        "       - Always includes PLANNER as base agent\n",
        "       - Sets initial priority structure\n",
        "       - Defines basic execution group\n",
        "\n",
        "    2. Response Analysis:\n",
        "       - Detects ReACT pattern markers\n",
        "       - Identifies agent requirements\n",
        "       - Determines execution priorities\n",
        "\n",
        "    3. Agent Coordination:\n",
        "       - Groups compatible agents for parallel execution\n",
        "       - Sets priority levels for sequential tasks\n",
        "       - Maintains execution dependencies\n",
        "    \"\"\"\n",
        "    try:\n",
        "        # Initialize with safe default configuration\n",
        "        analysis = {\n",
        "            \"required_agents\": [\"PLANNER\"],\n",
        "            \"priority\": {\"PLANNER\": 1},\n",
        "            \"concurrent_groups\": [[\"PLANNER\"]],\n",
        "            \"reasoning\": response\n",
        "        }\n",
        "\n",
        "        # Parse ReACT patterns for advanced coordination\n",
        "        if \"Thought:\" in response and \"Decision:\" in response:\n",
        "            # Check for NOTEWRITER requirements\n",
        "            if \"NOTEWRITER\" in response or \"note\" in response.lower():\n",
        "                analysis[\"required_agents\"].append(\"NOTEWRITER\")\n",
        "                analysis[\"priority\"][\"NOTEWRITER\"] = 2\n",
        "                # NOTEWRITER can run parallel with PLANNER\n",
        "                analysis[\"concurrent_groups\"] = [[\"PLANNER\", \"NOTEWRITER\"]]\n",
        "\n",
        "            # Check for ADVISOR requirements\n",
        "            if \"ADVISOR\" in response or \"guidance\" in response.lower():\n",
        "                analysis[\"required_agents\"].append(\"ADVISOR\")\n",
        "                analysis[\"priority\"][\"ADVISOR\"] = 3\n",
        "                # ADVISOR typically runs sequentially\n",
        "\n",
        "        return analysis\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"Parse error: {str(e)}\")\n",
        "        # Return safe default configuration\n",
        "        return {\n",
        "            \"required_agents\": [\"PLANNER\"],\n",
        "            \"priority\": {\"PLANNER\": 1},\n",
        "            \"concurrent_groups\": [[\"PLANNER\"]],\n",
        "            \"reasoning\": \"Fallback due to parse error\"\n",
        "        }"
      ],
      "metadata": {
        "id": "ZdQIcW1M26rS"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 👤 Profile Analyzer Agent"
      ],
      "metadata": {
        "id": "bNuQvLoIOMuF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "PROFILE_ANALYZER_PROMPT = \"\"\"You are a Profile Analysis Agent using the ReACT framework to analyze student profiles.\n",
        "\n",
        "    OBJECTIVE:\n",
        "    Analyze the student profile and extract key learning patterns that will impact their academic success.\n",
        "\n",
        "    REACT PATTERN:\n",
        "    Thought: Analyze what aspects of the profile need investigation\n",
        "    Action: Extract specific information from relevant profile sections\n",
        "    Observation: Note key patterns and implications\n",
        "    Response: Provide structured analysis\n",
        "\n",
        "    PROFILE DATA:\n",
        "    {profile}\n",
        "\n",
        "    ANALYSIS FRAMEWORK:\n",
        "    1. Learning Characteristics:\n",
        "        • Primary learning style\n",
        "        • Information processing patterns\n",
        "        • Attention span characteristics\n",
        "\n",
        "    2. Environmental Factors:\n",
        "        • Optimal study environment\n",
        "        • Distraction triggers\n",
        "        • Productive time periods\n",
        "\n",
        "    3. Executive Function:\n",
        "        • Task management patterns\n",
        "        • Focus duration limits\n",
        "        • Break requirements\n",
        "\n",
        "    4. Energy Management:\n",
        "        • Peak energy periods\n",
        "        • Recovery patterns\n",
        "        • Fatigue signals\n",
        "\n",
        "    INSTRUCTIONS:\n",
        "    1. Use the ReACT pattern for each analysis area\n",
        "    2. Provide specific, actionable observations\n",
        "    3. Note both strengths and challenges\n",
        "    4. Identify patterns that affect study planning\n",
        "\n",
        "    FORMAT YOUR RESPONSE AS:\n",
        "    Thought: [Initial analysis of profile components]\n",
        "    Action: [Specific areas being examined]\n",
        "    Observation: [Patterns and insights discovered]\n",
        "    Analysis Summary: [Structured breakdown of key findings]\n",
        "    Recommendations: [Specific adaptations needed]\n",
        "    \"\"\""
      ],
      "metadata": {
        "id": "HPU-eQMw0JMF"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "    Implementation Notes:\n",
        "    -------------------\n",
        "    1. Profile Analysis Process:\n",
        "       - Extracts profile data from state\n",
        "       - Applies ReACT framework for structured analysis\n",
        "       - Generates comprehensive learning insights\n",
        "       \n",
        "    2. ReACT Pattern Implementation:\n",
        "       The PROFILE_ANALYZER_PROMPT typically includes:\n",
        "       - Thought: Analysis of learning patterns and preferences\n",
        "       - Action: Identification of key learning traits\n",
        "       - Observation: Pattern recognition in academic history\n",
        "       - Decision: Synthesized learning profile recommendations\n",
        "       \n",
        "    3. LLM Integration:\n",
        "       - Uses structured prompting for consistent analysis\n",
        "       - Maintains conversation context through messages array\n",
        "       - Processes raw profile data through JSON serialization\n",
        "       \n",
        "    4. Result Structure:\n",
        "       Returns analysis in a format that:\n",
        "       - Can be combined with other agent outputs\n",
        "       - Provides clear learning preference insights\n",
        "       - Includes actionable recommendations"
      ],
      "metadata": {
        "id": "wRMJA2ie0zHv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "async def profile_analyzer(state: AcademicState) -> Dict:\n",
        "    \"\"\"\n",
        "    Analyzes student profile data to extract and interpret learning preferences using ReACT framework.\n",
        "\n",
        "    This agent specializes in:\n",
        "    1. Deep analysis of student learning profiles\n",
        "    2. Extraction of learning preferences and patterns\n",
        "    3. Interpretation of academic history and tendencies\n",
        "    4. Generation of personalized learning insights\n",
        "\n",
        "    Args:\n",
        "        state (AcademicState): Current academic state containing student profile data\n",
        "\n",
        "    Returns:\n",
        "        Dict: Structured analysis results including learning preferences and recommendations\n",
        "\n",
        "    \"\"\"\n",
        "    # Extract profile data from state\n",
        "    profile = state[\"profile\"]\n",
        "\n",
        "    # Assumes PROFILE_ANALYZER_PROMPT is defined elsewhere with ReACT framework structure\n",
        "    prompt = PROFILE_ANALYZER_PROMPT\n",
        "\n",
        "    # Construct message array for LLM interaction\n",
        "    messages = [\n",
        "        # System message defines analysis framework and expectations\n",
        "        {\"role\": \"system\", \"content\": prompt},\n",
        "        # User message contains serialized profile data for analysis\n",
        "        {\"role\": \"user\", \"content\": json.dumps(profile)}\n",
        "    ]\n",
        "\n",
        "    # Generate analysis using LLM\n",
        "    response = await llm.agenerate(messages)\n",
        "\n",
        "    # Format and structure the analysis results\n",
        "    return {\n",
        "        \"results\": {\n",
        "            \"profile_analysis\": {\n",
        "                \"analysis\": response  # Contains structured learning preference analysis\n",
        "            }\n",
        "        }\n",
        "    }"
      ],
      "metadata": {
        "id": "LJM4MK29OD4W"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 📅PlannerAgent"
      ],
      "metadata": {
        "id": "EWbQyHxODjaH"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "- Initialize PlannerAgent with Examples --> Create the Planning Workflow Graph and Return with compile the graph --> Creat a calendar Analysis and prompt -->  Create a plan analysis function and prompt -->  Create a planner generator function, define a ReACT prompt --> Execute the subgraph"
      ],
      "metadata": {
        "id": "JhcHrH9V6vEn"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class PlannerAgent(ReActAgent):\n",
        "    def __init__(self, llm):\n",
        "        super().__init__(llm)  # Initialize parent ReActAgent class\n",
        "        self.llm = llm\n",
        "        # Load example scenarios to help guide the AI's responses\n",
        "        self.few_shot_examples = self._initialize_fewshots()\n",
        "        # Create the workflow structure\n",
        "        self.workflow = self.create_subgraph()\n",
        "\n",
        "    def _initialize_fewshots(self):\n",
        "        \"\"\"\n",
        "        Define example scenarios to help the AI understand how to handle different situations\n",
        "        Each example shows:\n",
        "        - Input: The student's request\n",
        "        - Thought: The analysis process\n",
        "        - Action: What needs to be done\n",
        "        - Observation: What was found\n",
        "        - Plan: The detailed solution\n",
        "        \"\"\"\n",
        "        return [\n",
        "            {\n",
        "                \"input\": \"Help with exam prep while managing ADHD and football\",\n",
        "                \"thought\": \"Need to check calendar conflicts and energy patterns\",\n",
        "                \"action\": \"search_calendar\",\n",
        "                \"observation\": \"Football match at 6PM, exam tomorrow 9AM\",\n",
        "                \"plan\": \"\"\"ADHD-OPTIMIZED SCHEDULE:\n",
        "                    PRE-FOOTBALL (2PM-5PM):\n",
        "                    - 3x20min study sprints\n",
        "                    - Movement breaks\n",
        "                    - Quick rewards after each sprint\n",
        "\n",
        "                    FOOTBALL MATCH (6PM-8PM):\n",
        "                    - Use as dopamine reset\n",
        "                    - Formula review during breaks\n",
        "\n",
        "                    POST-MATCH (9PM-12AM):\n",
        "                    - Environment: Café noise\n",
        "                    - 15/5 study/break cycles\n",
        "                    - Location changes hourly\n",
        "\n",
        "                    EMERGENCY PROTOCOLS:\n",
        "                    - Focus lost → jumping jacks\n",
        "                    - Overwhelmed → room change\n",
        "                    - Brain fog → cold shower\"\"\"\n",
        "            },\n",
        "            {\n",
        "                \"input\": \"Struggling with multiple deadlines\",\n",
        "                \"thought\": \"Check task priorities and performance issues\",\n",
        "                \"action\": \"analyze_tasks\",\n",
        "                \"observation\": \"3 assignments due, lowest grade in Calculus\",\n",
        "                \"plan\": \"\"\"PRIORITY SCHEDULE:\n",
        "                    HIGH-FOCUS SLOTS:\n",
        "                    - Morning: Calculus practice\n",
        "                    - Post-workout: Assignments\n",
        "                    - Night: Quick reviews\n",
        "\n",
        "                    ADHD MANAGEMENT:\n",
        "                    - Task timer challenges\n",
        "                    - Reward system per completion\n",
        "                    - Study buddy accountability\"\"\"\n",
        "            }\n",
        "        ]\n",
        "    # Section 2: Create the Planning Workflow Graph and Return with compile the graph\n",
        "    def create_subgraph(self) -> StateGraph:\n",
        "        \"\"\"\n",
        "        Create a workflow graph that defines how the planner processes requests:\n",
        "        1. First analyzes calendar (calendar_analyzer)\n",
        "        2. Then analyzes tasks (task_analyzer)\n",
        "        3. Finally generates a plan (plan_generator)\n",
        "        \"\"\"\n",
        "        # Initialize a new graph using our AcademicState structure\n",
        "        subgraph = StateGraph(AcademicState)\n",
        "\n",
        "        # Add each processing step as a node in our graph\n",
        "        subgraph.add_node(\"calendar_analyzer\", self.calendar_analyzer)\n",
        "        subgraph.add_node(\"task_analyzer\", self.task_analyzer)\n",
        "        subgraph.add_node(\"plan_generator\", self.plan_generator)\n",
        "\n",
        "        # Connect the nodes in the order they should execute\n",
        "        subgraph.add_edge(\"calendar_analyzer\", \"task_analyzer\")\n",
        "        subgraph.add_edge(\"task_analyzer\", \"plan_generator\")\n",
        "\n",
        "        # Set where the workflow begins\n",
        "        subgraph.set_entry_point(\"calendar_analyzer\")\n",
        "\n",
        "        # Prepare the graph for use\n",
        "        return subgraph.compile()\n",
        "\n",
        "    async def calendar_analyzer(self, state: AcademicState) -> AcademicState:\n",
        "        \"\"\"\n",
        "        Analyze the student's calendar to find:\n",
        "        - Available study times\n",
        "        - Potential scheduling conflicts\n",
        "        - Energy patterns throughout the day\n",
        "        \"\"\"\n",
        "        # Get calendar events for the next 7 days\n",
        "        events = state[\"calendar\"].get(\"events\", [])\n",
        "        now = datetime.now(timezone.utc)\n",
        "        future = now + timedelta(days=7)\n",
        "\n",
        "        # Filter to only include upcoming events\n",
        "        filtered_events = [\n",
        "            event for event in events\n",
        "            if now <= datetime.fromisoformat(event[\"start\"][\"dateTime\"]) <= future\n",
        "        ]\n",
        "\n",
        "        # Create prompt for the AI to analyze the calendar\n",
        "        prompt = \"\"\"Analyze calendar events and identify:\n",
        "        Events: {events}\n",
        "\n",
        "        Focus on:\n",
        "        - Available time blocks\n",
        "        - Energy impact of activities\n",
        "        - Potential conflicts\n",
        "        - Recovery periods\n",
        "        - Study opportunity windows\n",
        "        - Activity patterns\n",
        "        - Schedule optimization\n",
        "        \"\"\"\n",
        "\n",
        "        # Ask AI to analyze the calendar\n",
        "        messages = [\n",
        "            {\"role\": \"system\", \"content\": prompt},\n",
        "            {\"role\": \"user\", \"content\": json.dumps(filtered_events)}\n",
        "        ]\n",
        "\n",
        "        response = await self.llm.agenerate(messages)\n",
        "        #cleaned_response = clean_llm_output({\"response\": response})\n",
        "\n",
        "        # Return the analysis results\n",
        "        return {\n",
        "            \"results\": {\n",
        "                \"calendar_analysis\": {\n",
        "                    \"analysis\":response\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "    async def task_analyzer(self, state: AcademicState) -> AcademicState:\n",
        "        \"\"\"\n",
        "        Analyze tasks to determine:\n",
        "        - Priority order\n",
        "        - Time needed for each task\n",
        "        - Best approach for completion\n",
        "        \"\"\"\n",
        "        tasks = state[\"tasks\"].get(\"tasks\", [])\n",
        "\n",
        "        # Create prompt for AI to analyze tasks\n",
        "        prompt = \"\"\"Analyze tasks and create priority structure:\n",
        "        Tasks: {tasks}\n",
        "\n",
        "        Consider:\n",
        "        - Urgency levels\n",
        "        - Task complexity\n",
        "        - Energy requirements\n",
        "        - Dependencies\n",
        "        - Required focus levels\n",
        "        - Time estimations\n",
        "        - Learning objectives\n",
        "        - Success criteria\n",
        "        \"\"\"\n",
        "\n",
        "        messages = [\n",
        "            {\"role\": \"system\", \"content\": prompt},\n",
        "            {\"role\": \"user\", \"content\": json.dumps(tasks)}\n",
        "        ]\n",
        "\n",
        "        response = await self.llm.agenerate(messages)\n",
        "        #cleaned_response = clean_llm_output({\"response\": response})\n",
        "\n",
        "        return {\n",
        "            \"results\": {\n",
        "                \"task_analysis\": {\n",
        "                    \"analysis\": response\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "\n",
        "    async def plan_generator(self, state: AcademicState) -> AcademicState:\n",
        "        \"\"\"\n",
        "        Create a comprehensive study plan by combining:\n",
        "        - Profile analysis (student's learning style)\n",
        "        - Calendar analysis (available time)\n",
        "        - Task analysis (what needs to be done)\n",
        "        \"\"\"\n",
        "        # Gather all previous analyses\n",
        "        profile_analysis = state[\"results\"][\"profile_analysis\"]\n",
        "        calendar_analysis = state[\"results\"][\"calendar_analysis\"]\n",
        "        task_analysis = state[\"results\"][\"task_analysis\"]\n",
        "\n",
        "        # Create detailed prompt for AI to generate plan\n",
        "        prompt = f\"\"\"AI Planning Assistant: Create focused study plan using ReACT framework.\n",
        "\n",
        "          INPUT CONTEXT:\n",
        "          - Profile Analysis: {profile_analysis}\n",
        "          - Calendar Analysis: {calendar_analysis}\n",
        "          - Task Analysis: {task_analysis}\n",
        "\n",
        "          EXAMPLES:\n",
        "          {json.dumps(self.few_shot_examples, indent=2)}\n",
        "\n",
        "          INSTRUCTIONS:\n",
        "          1. Follow ReACT pattern:\n",
        "            Thought: Analyze situation and needs\n",
        "            Action: Consider all analyses\n",
        "            Observation: Synthesize findings\n",
        "            Plan: Create structured plan\n",
        "\n",
        "          2. Address:\n",
        "            - ADHD management strategies\n",
        "            - Energy level optimization\n",
        "            - Task chunking methods\n",
        "            - Focus period scheduling\n",
        "            - Environment switching tactics\n",
        "            - Recovery period planning\n",
        "            - Social/sport activity balance\n",
        "\n",
        "          3. Include:\n",
        "            - Emergency protocols\n",
        "            - Backup strategies\n",
        "            - Quick wins\n",
        "            - Reward system\n",
        "            - Progress tracking\n",
        "            - Adjustment triggers\n",
        "\n",
        "          Pls act as an intelligent tool to help the students reach their goals or overcome struggles and answer with informal words.\n",
        "\n",
        "          FORMAT:\n",
        "          Thought: [reasoning and situation analysis]\n",
        "          Action: [synthesis approach]\n",
        "          Observation: [key findings]\n",
        "          Plan: [actionable steps and structural schedule]\n",
        "          \"\"\"\n",
        "\n",
        "\n",
        "        messages = [\n",
        "            {\"role\": \"system\", \"content\": prompt},\n",
        "            {\"role\": \"user\", \"content\": state[\"messages\"][-1].content}\n",
        "        ]\n",
        "        # temperature is like a randomness of LLM response, 0.5 is in the middle\n",
        "        response = await self.llm.agenerate(messages, temperature=0.5)\n",
        "\n",
        "        # Clean the response before returning\n",
        "        #cleaned_response = clean_llm_output({\"response\": response})\n",
        "\n",
        "        return {\n",
        "                \"results\": {\n",
        "                    \"final_plan\": {\n",
        "                        \"plan\": response\n",
        "                    }\n",
        "                }\n",
        "        }\n",
        "\n",
        "    async def __call__(self, state: AcademicState) -> Dict:\n",
        "        \"\"\"\n",
        "        Main execution method that runs the entire planning workflow:\n",
        "        1. Analyze calendar\n",
        "        2. Analyze tasks\n",
        "        3. Generate plan\n",
        "        \"\"\"\n",
        "        try:\n",
        "            final_state = await self.workflow.ainvoke(state)\n",
        "            # Clean the generated notes before returning\n",
        "            notes = final_state[\"results\"].get(\"generated_notes\", {})\n",
        "            #cleaned_notes = clean_llm_output({\"notes\": notes})\n",
        "            return {\"notes\": final_state[\"results\"].get(\"generated_notes\")}\n",
        "            #return {\"notes\": cleaned_notes.get(\"notes\")}\n",
        "        except Exception as e:\n",
        "            return {\"notes\": \"Error generating notes. Please try again.\"}\n",
        "\n"
      ],
      "metadata": {
        "id": "mRAAYta16WT8"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#📚 NoteWriterAgent"
      ],
      "metadata": {
        "id": "_Hkj0YUqD9y_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class NoteWriterAgent(ReActAgent):\n",
        "    \"\"\"NoteWriter agent with its own subgraph workflow for note generation.\n",
        "    This agent specializes in creating personalized study materials by analyzing\n",
        "    learning styles and generating structured notes.\"\"\"\n",
        "\n",
        "    def __init__(self, llm):\n",
        "        \"\"\"Initialize the NoteWriter agent with an LLM backend and example templates.\n",
        "\n",
        "        Args:\n",
        "            llm: Language model instance for text generation\n",
        "        \"\"\"\n",
        "        super().__init__(llm)\n",
        "        self.llm = llm\n",
        "        self.few_shot_examples = [\n",
        "            {\n",
        "                \"input\": \"Need to cram Calculus III for tomorrow\",\n",
        "                \"template\": \"Quick Review\",\n",
        "                \"notes\": \"\"\"CALCULUS III ESSENTIALS:\n",
        "\n",
        "                1. CORE CONCEPTS (80/20 Rule):\n",
        "                   • Multiple Integrals → volume/area\n",
        "                   • Vector Calculus → flow/force/rotation\n",
        "                   • KEY FORMULAS:\n",
        "                     - Triple integrals in cylindrical/spherical coords\n",
        "                     - Curl, divergence, gradient relationships\n",
        "\n",
        "                2. COMMON EXAM PATTERNS:\n",
        "                   • Find critical points\n",
        "                   • Calculate flux/work\n",
        "                   • Optimize with constraints\n",
        "\n",
        "                3. QUICKSTART GUIDE:\n",
        "                   • Always draw 3D diagrams\n",
        "                   • Check units match\n",
        "                   • Use symmetry to simplify\n",
        "\n",
        "                4. EMERGENCY TIPS:\n",
        "                   • If stuck, try converting coordinates\n",
        "                   • Check boundary conditions\n",
        "                   • Look for special patterns\"\"\"\n",
        "            }\n",
        "        ]\n",
        "        self.workflow = self.create_subgraph()\n",
        "\n",
        "    def create_subgraph(self) -> StateGraph:\n",
        "        \"\"\"Creates NoteWriter's internal workflow as a state machine.\n",
        "\n",
        "        The workflow consists of two main steps:\n",
        "        1. Analyze learning style and content requirements\n",
        "        2. Generate personalized notes\n",
        "\n",
        "        Returns:\n",
        "            StateGraph: Compiled workflow graph\n",
        "        \"\"\"\n",
        "        subgraph = StateGraph(AcademicState)\n",
        "\n",
        "        # Define the core workflow nodes\n",
        "        subgraph.add_node(\"notewriter_analyze\", self.analyze_learning_style)\n",
        "        subgraph.add_node(\"notewriter_generate\", self.generate_notes)\n",
        "\n",
        "        # Create the workflow sequence:\n",
        "        # START -> analyze -> generate -> END\n",
        "        subgraph.add_edge(START, \"notewriter_analyze\")\n",
        "        subgraph.add_edge(\"notewriter_analyze\", \"notewriter_generate\")\n",
        "        subgraph.add_edge(\"notewriter_generate\", END)\n",
        "\n",
        "        return subgraph.compile()\n",
        "\n",
        "    async def analyze_learning_style(self, state: AcademicState) -> AcademicState:\n",
        "        \"\"\"Analyzes student profile and request to determine optimal note structure.\n",
        "\n",
        "        Uses the LLM to analyze:\n",
        "        - Student's learning style preferences\n",
        "        - Specific content request\n",
        "        - Time constraints and requirements\n",
        "\n",
        "        Args:\n",
        "            state: Current academic state containing student profile and messages\n",
        "\n",
        "        Returns:\n",
        "            Updated state with learning analysis results\n",
        "        \"\"\"\n",
        "        profile = state[\"profile\"]\n",
        "        learning_style = profile[\"learning_preferences\"][\"learning_style\"]\n",
        "        # Construct analysis prompt with specific formatting requirements\n",
        "\n",
        "        prompt = f\"\"\"Analyze content requirements and determine optimal note structure:\n",
        "\n",
        "        STUDENT PROFILE:\n",
        "        - Learning Style: {json.dumps(learning_style, indent=2)}\n",
        "        - Request: {state['messages'][-1].content}\n",
        "\n",
        "        FORMAT:\n",
        "        1. Key Topics (80/20 principle)\n",
        "        2. Learning Style Adaptations\n",
        "        3. Time Management Strategy\n",
        "        4. Quick Reference Format\n",
        "\n",
        "        FOCUS ON:\n",
        "        - Essential concepts that give maximum understanding\n",
        "        - Visual and interactive elements\n",
        "        - Time-optimized study methods\n",
        "        \"\"\"\n",
        "\n",
        "        response = await self.llm.agenerate([\n",
        "            {\"role\": \"system\", \"content\": prompt}\n",
        "        ])\n",
        "        #cleaned_response = clean_llm_output({\"response\": response})\n",
        "\n",
        "\n",
        "        return {\n",
        "            \"results\": {\n",
        "                \"learning_analysis\": {\n",
        "                    \"analysis\": response\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "\n",
        "    async def generate_notes(self, state: AcademicState) -> AcademicState:\n",
        "        \"\"\"Generates personalized study notes based on the learning analysis.\n",
        "\n",
        "        Uses the LLM to create structured notes that are:\n",
        "        - Adapted to the student's learning style\n",
        "        - Focused on essential concepts (80/20 principle)\n",
        "        - Time-optimized for the study period\n",
        "\n",
        "        Args:\n",
        "            state: Current academic state with learning analysis\n",
        "        Returns:\n",
        "            Updated state with generated notes\n",
        "        \"\"\"\n",
        "\n",
        "        analysis = state[\"results\"].get(\"learning_analysis\", \"\")\n",
        "        learning_style = state[\"profile\"][\"learning_preferences\"][\"learning_style\"]\n",
        "\n",
        "        # Build prompt using analysis and few-shot examples\n",
        "        prompt = f\"\"\"Create concise, high-impact study materials based on analysis:\n",
        "\n",
        "        ANALYSIS: {analysis}\n",
        "        LEARNING STYLE: {json.dumps(learning_style, indent=2)}\n",
        "        REQUEST: {state['messages'][-1].content}\n",
        "\n",
        "        EXAMPLES:\n",
        "        {json.dumps(self.few_shot_examples, indent=2)}\n",
        "\n",
        "        FORMAT:\n",
        "        **THREE-WEEK INTENSIVE STUDY PLANNER**\n",
        "\n",
        "        [Generate structured notes with:]\n",
        "        1. Weekly breakdown\n",
        "        2. Daily focus areas\n",
        "        3. Core concepts\n",
        "        4. Emergency tips\n",
        "        \"\"\"\n",
        "\n",
        "        response = await self.llm.agenerate([\n",
        "            {\"role\": \"system\", \"content\": prompt}\n",
        "        ])\n",
        "        #cleaned_response = clean_llm_output({\"response\": response})\n",
        "        # if \"results\" not in state:\n",
        "        #     state[\"results\"] = {}\n",
        "        # state[\"results\"][\"generated_notes\"] = response\n",
        "        # return state\n",
        "        return {\n",
        "            \"results\": {\n",
        "                \"generated_notes\": {\n",
        "                    \"notes\": response\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "\n",
        "\n",
        "    async def __call__(self, state: AcademicState) -> Dict:\n",
        "        \"\"\"Main execution method for the NoteWriter agent.\n",
        "\n",
        "        Executes the complete workflow:\n",
        "        1. Analyzes learning requirements\n",
        "        2. Generates personalized notes\n",
        "        3. Cleans and returns the results\n",
        "\n",
        "        Args:\n",
        "            state: Initial academic state\n",
        "\n",
        "        Returns:\n",
        "            Dict containing generated notes or error message\n",
        "        \"\"\"\n",
        "        try:\n",
        "            final_state = await self.workflow.ainvoke(state)\n",
        "            # Clean the generated notes before returning\n",
        "            notes = final_state[\"results\"].get(\"generated_notes\", {})\n",
        "            #cleaned_notes = clean_llm_output({\"notes\": notes})\n",
        "            return {\"notes\": final_state[\"results\"].get(\"generated_notes\")}\n",
        "            #return {\"notes\": cleaned_notes.get(\"notes\")}\n",
        "        except Exception as e:\n",
        "            return {\"notes\": \"Error generating notes. Please try again.\"}"
      ],
      "metadata": {
        "id": "K5HFsCgu3OOm"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 👩🏼‍🏫 AdvisorAgent"
      ],
      "metadata": {
        "id": "R3LnR-SlEAhz"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class AdvisorAgent(ReActAgent):\n",
        "    \"\"\"Academic advisor agent with subgraph workflow for personalized guidance.\n",
        "    This agent specializes in analyzing student situations and providing\n",
        "    tailored academic advice considering learning styles and time constraints.\"\"\"\n",
        "\n",
        "    def __init__(self, llm):\n",
        "        \"\"\"Initialize the Advisor agent with an LLM backend and example templates.\n",
        "\n",
        "        Args:\n",
        "            llm: Language model instance for text generation\n",
        "        \"\"\"\n",
        "        super().__init__(llm)\n",
        "        self.llm = llm\n",
        "\n",
        "        # Define comprehensive examples for guidance generation\n",
        "        # These examples help the LLM understand the expected format and depth\n",
        "        self.few_shot_examples = [\n",
        "            {\n",
        "                \"request\": \"Managing multiple deadlines with limited time\",\n",
        "                \"profile\": {\n",
        "                    \"learning_style\": \"visual\",\n",
        "                    \"workload\": \"heavy\",\n",
        "                    \"time_constraints\": [\"2 hackathons\", \"project\", \"exam\"]\n",
        "                },\n",
        "                \"advice\": \"\"\"PRIORITY-BASED SCHEDULE:\n",
        "\n",
        "                1. IMMEDIATE ACTIONS\n",
        "                   • Create visual timeline of all deadlines\n",
        "                   • Break each task into 45-min chunks\n",
        "                   • Schedule high-focus work in mornings\n",
        "\n",
        "                2. WORKLOAD MANAGEMENT\n",
        "                   • Hackathons: Form team early, set clear roles\n",
        "                   • Project: Daily 2-hour focused sessions\n",
        "                   • Exam: Interleaved practice with breaks\n",
        "\n",
        "                3. ENERGY OPTIMIZATION\n",
        "                   • Use Pomodoro (25/5) for intensive tasks\n",
        "                   • Physical activity between study blocks\n",
        "                   • Regular progress tracking\n",
        "\n",
        "                4. EMERGENCY PROTOCOLS\n",
        "                   • If overwhelmed: Take 10min reset break\n",
        "                   • If stuck: Switch tasks or environments\n",
        "                   • If tired: Quick power nap, then review\"\"\"\n",
        "            }\n",
        "        ]\n",
        "        # Initialize the agent's workflow state machine\n",
        "        self.workflow = self.create_subgraph()\n",
        "\n",
        "    def create_subgraph(self) -> StateGraph:\n",
        "      \"\"\"Creates Advisor's internal workflow as a state machine.\n",
        "\n",
        "        The workflow consists of two main stages:\n",
        "        1. Situation analysis - Understanding student needs\n",
        "        2. Guidance generation - Creating personalized advice\n",
        "\n",
        "        Returns:\n",
        "            StateGraph: Compiled workflow graph\n",
        "      \"\"\"\n",
        "      subgraph = StateGraph(AcademicState)\n",
        "\n",
        "      # Add nodes for analysis and guidance - use consistent names\n",
        "      subgraph.add_node(\"advisor_analyze\", self.analyze_situation)\n",
        "      subgraph.add_node(\"advisor_generate\", self.generate_guidance)\n",
        "\n",
        "      # Connect workflow - use the new node names\n",
        "      subgraph.add_edge(START, \"advisor_analyze\")\n",
        "      subgraph.add_edge(\"advisor_analyze\", \"advisor_generate\")\n",
        "      subgraph.add_edge(\"advisor_generate\", END)\n",
        "\n",
        "      return subgraph.compile()\n",
        "\n",
        "    async def analyze_situation(self, state: AcademicState) -> AcademicState:\n",
        "        \"\"\"Analyzes student's current academic situation and needs.\n",
        "\n",
        "        Evaluates:\n",
        "        - Student profile and preferences\n",
        "        - Current challenges and constraints\n",
        "        - Learning style compatibility\n",
        "        - Time and stress management needs\n",
        "\n",
        "        Args:\n",
        "            state: Current academic state with student profile and request\n",
        "\n",
        "        Returns:\n",
        "            Updated state with situation analysis results\n",
        "        \"\"\"\n",
        "        profile = state[\"profile\"]\n",
        "        learning_prefs = profile.get(\"learning_preferences\", {})\n",
        "\n",
        "        prompt = f\"\"\"Analyze student situation and determine guidance approach:\n",
        "\n",
        "        CONTEXT:\n",
        "        - Profile: {json.dumps(profile, indent=2)}\n",
        "        - Learning Preferences: {json.dumps(learning_prefs, indent=2)}\n",
        "        - Request: {state['messages'][-1].content}\n",
        "\n",
        "        ANALYZE:\n",
        "        1. Current challenges\n",
        "        2. Learning style compatibility\n",
        "        3. Time management needs\n",
        "        4. Stress management requirements\n",
        "        \"\"\"\n",
        "\n",
        "        response = await self.llm.agenerate([\n",
        "            {\"role\": \"system\", \"content\": prompt}\n",
        "        ])\n",
        "\n",
        "        # if \"results\" not in state:\n",
        "        #     state[\"results\"] = {}\n",
        "        # state[\"results\"][\"situation_analysis\"] = response\n",
        "        # return state\n",
        "        #Clean the response before returning\n",
        "        #cleaned_response = clean_llm_output({\"response\": response})\n",
        "\n",
        "        return {\n",
        "            \"results\": {\n",
        "                \"situation_analysis\": {\n",
        "                    \"analysis\": response\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "\n",
        "    async def generate_guidance(self, state: AcademicState) -> AcademicState:\n",
        "        \"\"\"Generates personalized academic guidance based on situation analysis.\n",
        "\n",
        "        Creates structured advice focusing on:\n",
        "        - Immediate actionable steps\n",
        "        - Schedule optimization\n",
        "        - Energy and resource management\n",
        "        - Support strategies\n",
        "        - Contingency planning\n",
        "\n",
        "        Args:\n",
        "            state: Current academic state with situation analysis\n",
        "\n",
        "        Returns:\n",
        "            Updated state with generated guidance\n",
        "        \"\"\"\n",
        "\n",
        "        analysis = state[\"results\"].get(\"situation_analysis\", \"\")\n",
        "\n",
        "        prompt = f\"\"\"Generate personalized academic guidance based on analysis:\n",
        "\n",
        "        ANALYSIS: {analysis}\n",
        "        EXAMPLES: {json.dumps(self.few_shot_examples, indent=2)}\n",
        "\n",
        "        FORMAT:\n",
        "        1. Immediate Action Steps\n",
        "        2. Schedule Optimization\n",
        "        3. Energy Management\n",
        "        4. Support Strategies\n",
        "        5. Emergency Protocols\n",
        "        \"\"\"\n",
        "\n",
        "        response = await self.llm.agenerate([\n",
        "            {\"role\": \"system\", \"content\": prompt}\n",
        "        ])\n",
        "\n",
        "        # if \"results\" not in state:\n",
        "        #     state[\"results\"] = {}\n",
        "        # state[\"results\"][\"guidance\"] = response\n",
        "        # return state\n",
        "        #cleaned_response = clean_llm_output({\"response\": response})\n",
        "\n",
        "        return {\n",
        "            \"results\": {\n",
        "                \"guidance\": {\n",
        "                    \"advice\": response\n",
        "                }\n",
        "            }\n",
        "        }\n",
        "\n",
        "    async def __call__(self, state: AcademicState) -> Dict:\n",
        "        \"\"\"Main execution method for the Advisor agent.\n",
        "\n",
        "        Executes the complete advisory workflow:\n",
        "        1. Analyzes student situation\n",
        "        2. Generates personalized guidance\n",
        "        3. Returns formatted results with metadata\n",
        "\n",
        "        Args:\n",
        "            state: Initial academic state\n",
        "\n",
        "        Returns:\n",
        "            Dict containing guidance and metadata or error message\n",
        "\n",
        "        Note:\n",
        "            Includes metadata about guidance specificity and learning style consideration\n",
        "        \"\"\"\n",
        "\n",
        "        try:\n",
        "            final_state = await self.workflow.ainvoke(state)\n",
        "            return {\n",
        "                \"advisor_output\": {\n",
        "                    \"guidance\": final_state[\"results\"].get(\"guidance\"),\n",
        "                    \"metadata\": {\n",
        "                        \"course_specific\": True,\n",
        "                        \"considers_learning_style\": True\n",
        "                    }\n",
        "                }\n",
        "            }\n",
        "        except Exception as e:\n",
        "            return {\n",
        "                \"advisor_output\": {\n",
        "                    \"guidance\": \"Error generating guidance. Please try again.\"\n",
        "                }\n",
        "            }"
      ],
      "metadata": {
        "id": "p_slfG6D3TwW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 🚀Multi-Agent Workflow Orchestration and State Graph Construction"
      ],
      "metadata": {
        "id": "IcBmQIHTECsg"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "- This code demonstrates how to create a coordinated workflow system (StateGraph) that manages multiple AI academic support agents running in parallel.\n",
        "- Key Components:\n",
        "    - State Graph Construction\n",
        "    - Building a workflow using nodes and edges\n",
        "    - Defining execution paths between agents\n",
        "    - Managing state transitions\n",
        "    - Parallel Agent Coordination\n",
        "\n",
        "- 3 main agents working together:\n",
        "    - PlannerAgent (scheduling/calendar)\n",
        "    - NoteWriterAgent (study materials)\n",
        "    - AdvisorAgent (academic guidance)"
      ],
      "metadata": {
        "id": "gFIkfipJG1KP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "- Orchestrator: Coordinates multiple agents' workflows\n",
        "- Router: Directs requests to appropriate agents\n",
        "- State Manager: Maintains workflow state and transitions\n",
        "- Completion Handler: Determines when all required work is done"
      ],
      "metadata": {
        "id": "2VDHAJ77H1xs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def create_agents_graph(llm) -> StateGraph:\n",
        "    \"\"\"Creates a coordinated workflow graph for multiple AI agents.\n",
        "\n",
        "    This orchestration system manages parallel execution of three specialized agents:\n",
        "    - PlannerAgent: Handles scheduling and calendar management\n",
        "    - NoteWriterAgent: Creates personalized study materials\n",
        "    - AdvisorAgent: Provides academic guidance and support\n",
        "\n",
        "    The workflow uses a state machine approach with conditional routing based on\n",
        "    analysis of student needs.\n",
        "\n",
        "    Args:\n",
        "        llm: Language model instance shared across all agents\n",
        "\n",
        "    Returns:\n",
        "        StateGraph: Compiled workflow graph with parallel execution paths\n",
        "    \"\"\"\n",
        "    # Initialize main workflow state machine\n",
        "    workflow = StateGraph(AcademicState)\n",
        "\n",
        "    # Create instances of our specialized agents\n",
        "    # Each agent has its own subgraph for internal operations\n",
        "    planner_agent = PlannerAgent(llm)\n",
        "    notewriter_agent = NoteWriterAgent(llm)\n",
        "    advisor_agent = AdvisorAgent(llm)\n",
        "    executor = AgentExecutor(llm)\n",
        "\n",
        "    # === MAIN WORKFLOW NODES ===\n",
        "    # These nodes handle high-level coordination and analysis\n",
        "    workflow.add_node(\"coordinator\", coordinator_agent)  # Initial request analysis\n",
        "    workflow.add_node(\"profile_analyzer\", profile_analyzer)  # Student profile analysis\n",
        "    workflow.add_node(\"execute\", executor.execute)  # Final execution node\n",
        "\n",
        "    # === PARALLEL EXECUTION ROUTING ===\n",
        "    def route_to_parallel_agents(state: AcademicState) -> List[str]:\n",
        "        \"\"\"Determines which agents should process the current request.\n",
        "\n",
        "        Analyzes coordinator's output to route work to appropriate agents.\n",
        "        Defaults to planner if no specific agents are required.\n",
        "\n",
        "        Args:\n",
        "            state: Current academic state with coordinator analysis\n",
        "\n",
        "        Returns:\n",
        "            List of next node names to execute\n",
        "        \"\"\"\n",
        "        analysis = state[\"results\"].get(\"coordinator_analysis\", {})\n",
        "        required_agents = analysis.get(\"required_agents\", [])\n",
        "        next_nodes = []\n",
        "\n",
        "        # Route to appropriate agent entry points based on analysis\n",
        "        if \"PLANNER\" in required_agents:\n",
        "            next_nodes.append(\"calendar_analyzer\")\n",
        "        if \"NOTEWRITER\" in required_agents:\n",
        "            next_nodes.append(\"notewriter_analyze\")\n",
        "        if \"ADVISOR\" in required_agents:\n",
        "            next_nodes.append(\"advisor_analyze\")\n",
        "\n",
        "        # Default to planner if no specific agents requested\n",
        "        return next_nodes if next_nodes else [\"calendar_analyzer\"]\n",
        "\n",
        "    # === AGENT SUBGRAPH NODES ===\n",
        "    # Add nodes for Planner agent's workflow\n",
        "    workflow.add_node(\"calendar_analyzer\", planner_agent.calendar_analyzer)\n",
        "    workflow.add_node(\"task_analyzer\", planner_agent.task_analyzer)\n",
        "    workflow.add_node(\"plan_generator\", planner_agent.plan_generator)\n",
        "\n",
        "    # Add nodes for NoteWriter agent's workflow\n",
        "    workflow.add_node(\"notewriter_analyze\", notewriter_agent.analyze_learning_style)\n",
        "    workflow.add_node(\"notewriter_generate\", notewriter_agent.generate_notes)\n",
        "\n",
        "    # Add nodes for Advisor agent's workflow\n",
        "    workflow.add_node(\"advisor_analyze\", advisor_agent.analyze_situation)\n",
        "    workflow.add_node(\"advisor_generate\", advisor_agent.generate_guidance)\n",
        "\n",
        "    # === WORKFLOW CONNECTIONS ===\n",
        "    # Main workflow entry\n",
        "    workflow.add_edge(START, \"coordinator\")\n",
        "    workflow.add_edge(\"coordinator\", \"profile_analyzer\")\n",
        "\n",
        "    # Connect profile analyzer to potential parallel paths\n",
        "    workflow.add_conditional_edges(\n",
        "        \"profile_analyzer\",\n",
        "        route_to_parallel_agents,\n",
        "        [\"calendar_analyzer\", \"notewriter_analyze\", \"advisor_analyze\"]\n",
        "    )\n",
        "\n",
        "    # Connect Planner agent's internal workflow\n",
        "    workflow.add_edge(\"calendar_analyzer\", \"task_analyzer\")\n",
        "    workflow.add_edge(\"task_analyzer\", \"plan_generator\")\n",
        "    workflow.add_edge(\"plan_generator\", \"execute\")\n",
        "\n",
        "    # Connect NoteWriter agent's internal workflow\n",
        "    workflow.add_edge(\"notewriter_analyze\", \"notewriter_generate\")\n",
        "    workflow.add_edge(\"notewriter_generate\", \"execute\")\n",
        "\n",
        "    # Connect Advisor agent's internal workflow\n",
        "    workflow.add_edge(\"advisor_analyze\", \"advisor_generate\")\n",
        "    workflow.add_edge(\"advisor_generate\", \"execute\")\n",
        "\n",
        "    # === WORKFLOW COMPLETION CHECKING ===\n",
        "    def should_end(state) -> Union[Literal[\"coordinator\"], Literal[END]]:\n",
        "        \"\"\"Determines if all required agents have completed their tasks.\n",
        "\n",
        "        Compares the set of completed agent outputs against required agents\n",
        "        to decide whether to end or continue the workflow.\n",
        "\n",
        "        Args:\n",
        "            state: Current academic state\n",
        "\n",
        "        Returns:\n",
        "            Either \"coordinator\" to continue or END to finish\n",
        "        \"\"\"\n",
        "        analysis = state[\"results\"].get(\"coordinator_analysis\", {})\n",
        "        executed = set(state[\"results\"].get(\"agent_outputs\", {}).keys())\n",
        "        required = set(a.lower() for a in analysis.get(\"required_agents\", []))\n",
        "        return END if required.issubset(executed) else \"coordinator\"\n",
        "\n",
        "    # Add conditional loop back to coordinator if needed\n",
        "    workflow.add_conditional_edges(\n",
        "        \"execute\",\n",
        "        should_end,\n",
        "        {\n",
        "            \"coordinator\": \"coordinator\",  # Loop back if more work needed\n",
        "            END: END  # End workflow if all agents complete\n",
        "        }\n",
        "    )\n",
        "\n",
        "    # Compile and return the complete workflow\n",
        "    return workflow.compile()"
      ],
      "metadata": {
        "id": "sct5u5kYFLMZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Run Streamlined Output"
      ],
      "metadata": {
        "id": "RDKtb_MkULJV"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "async def run_all_system(profile_json: str, calendar_json: str, task_json: str):\n",
        "    \"\"\"Run the entire academic assistance system with improved output handling.\n",
        "\n",
        "    This is the main entry point for the ATLAS (Academic Task Learning Agent System).\n",
        "    It handles initialization, user interaction, workflow execution, and result presentation.\n",
        "\n",
        "    Args:\n",
        "        profile_json: JSON string containing student profile data\n",
        "        calendar_json: JSON string containing calendar/schedule data\n",
        "        task_json: JSON string containing academic tasks data\n",
        "\n",
        "    Returns:\n",
        "        Tuple[Dict, Dict]: Coordinator output and final state, or (None, None) on error\n",
        "\n",
        "    Features:\n",
        "        - Rich console interface with status updates\n",
        "        - Async streaming of workflow steps\n",
        "        - Comprehensive error handling\n",
        "        - Live progress feedback\n",
        "    \"\"\"\n",
        "    try:\n",
        "        # Initialize rich console for enhanced UI\n",
        "        console = Console()\n",
        "\n",
        "        # Display welcome banner\n",
        "        console.print(\"\\n[bold magenta]🎓 ATLAS: Academic Task Learning Agent System[/bold magenta]\")\n",
        "        console.print(\"[italic blue]Initializing academic support system...[/italic blue]\\n\")\n",
        "\n",
        "        # Initialize core system components\n",
        "        # NeMoLLaMa is the language model backend\n",
        "        llm = NeMoLLaMa(NEMOTRON_4_340B_INSTRUCT_KEY)\n",
        "\n",
        "        # DataManager handles all data loading and access\n",
        "        dm = DataManager()\n",
        "        dm.load_data(profile_json, calendar_json, task_json)\n",
        "\n",
        "        # Get user request\n",
        "        console.print(\"[bold green]Please enter your academic request:[/bold green]\")\n",
        "        user_input = str(input())\n",
        "        console.print(f\"\\n[dim italic]Processing request: {user_input}[/dim italic]\\n\")\n",
        "\n",
        "        # Construct initial state object\n",
        "        # This contains all context needed by the agents\n",
        "        state = {\n",
        "            \"messages\": [HumanMessage(content=user_input)],  # User request\n",
        "            \"profile\": dm.get_student_profile(\"student_123\"),  # Student info\n",
        "            \"calendar\": {\"events\": dm.get_upcoming_events()},  # Schedule\n",
        "            \"tasks\": {\"tasks\": dm.get_active_tasks()},        # Active tasks\n",
        "            \"results\": {}                                     # Will store agent outputs\n",
        "        }\n",
        "\n",
        "        # Initialize workflow graph for agent orchestration\n",
        "        graph = create_agents_graph(llm)\n",
        "\n",
        "        console.print(\"[bold cyan]System initialized and processing request...[/bold cyan]\\n\")\n",
        "        # Add visualization here\n",
        "        console.print(\"[bold cyan]Workflow Graph Structure:[/bold cyan]\\n\")\n",
        "        display(Image(graph.get_graph().draw_mermaid_png()))\n",
        "\n",
        "        # Track important state transitions\n",
        "        coordinator_output = None  # Initial analysis\n",
        "        final_state = None        # Final results\n",
        "\n",
        "        # Process workflow with live status updates\n",
        "        with console.status(\"[bold green]Processing...\", spinner=\"dots\") as status:\n",
        "            # Stream workflow steps asynchronously\n",
        "            async for step in graph.astream(state):\n",
        "                # Capture coordinator analysis when available\n",
        "                if \"coordinator_analysis\" in step.get(\"results\", {}):\n",
        "                    coordinator_output = step\n",
        "                    analysis = coordinator_output[\"results\"][\"coordinator_analysis\"]\n",
        "\n",
        "                    # Display selected agents for transparency\n",
        "                    console.print(\"\\n[bold cyan]Selected Agents:[/bold cyan]\")\n",
        "                    for agent in analysis.get(\"required_agents\", []):\n",
        "                        console.print(f\"• {agent}\")\n",
        "\n",
        "                # Capture final execution state\n",
        "                if \"execute\" in step:\n",
        "                    final_state = step\n",
        "\n",
        "        # # Display formatted results if available\n",
        "        # if final_state:\n",
        "        #     display_formatted_output(final_state)\n",
        "        # Replace with simpler console output:\n",
        "        if final_state:\n",
        "            agent_outputs = final_state.get(\"execute\", {}).get(\"results\", {}).get(\"agent_outputs\", {})\n",
        "\n",
        "            # Simple console output for each agent\n",
        "            for agent, output in agent_outputs.items():\n",
        "                console.print(f\"\\n[bold cyan]{agent.upper()} Output:[/bold cyan]\")\n",
        "\n",
        "                # Handle nested dictionary output\n",
        "                if isinstance(output, dict):\n",
        "                    for key, value in output.items():\n",
        "                        if isinstance(value, dict):\n",
        "                            for subkey, subvalue in value.items():\n",
        "                                if subvalue and isinstance(subvalue, str):\n",
        "                                    console.print(subvalue.strip())\n",
        "                        elif value and isinstance(value, str):\n",
        "                            console.print(value.strip())\n",
        "                # Handle direct string output\n",
        "                elif isinstance(output, str):\n",
        "                    console.print(output.strip())\n",
        "\n",
        "        # Indicate completion\n",
        "        console.print(\"\\n[bold green]✓[/bold green] [bold]Task completed![/bold]\")\n",
        "        return coordinator_output, final_state\n",
        "\n",
        "    except Exception as e:\n",
        "        # Comprehensive error handling with stack trace\n",
        "        console.print(f\"\\n[bold red]System error:[/bold red] {str(e)}\")\n",
        "        console.print(\"[yellow]Stack trace:[/yellow]\")\n",
        "        import traceback\n",
        "        console.print(traceback.format_exc())\n",
        "        return None, None"
      ],
      "metadata": {
        "id": "LGSWX-KPpTEA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# 🚀Upload 3 tasks, events and profile samples and run the system"
      ],
      "metadata": {
        "id": "_zH88JfdBhH_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "async def load_json_and_test():\n",
        "    \"\"\"Load JSON files and run the academic assistance system.\"\"\"\n",
        "    print(\"Academic Assistant Test Setup\")\n",
        "    print(\"-\" * 50)\n",
        "    print(\"\\nPlease upload your JSON files...\")\n",
        "\n",
        "    try:\n",
        "        # Handle file upload\n",
        "        uploaded = files.upload()\n",
        "        if not uploaded:\n",
        "            print(\"No files were uploaded.\")\n",
        "            return\n",
        "\n",
        "        # Define patterns for matching file types\n",
        "        patterns = {\n",
        "            'profile': r'profile.*\\.json$',\n",
        "            'calendar': r'calendar.*\\.json$',\n",
        "            'task': r'task.*\\.json$'\n",
        "        }\n",
        "\n",
        "        # Find matching files\n",
        "        found_files = {\n",
        "            file_type: next((\n",
        "                f for f in uploaded.keys()\n",
        "                if re.match(pattern, f, re.IGNORECASE)\n",
        "            ), None)\n",
        "            for file_type, pattern in patterns.items()\n",
        "        }\n",
        "\n",
        "        # Check if all required files are present\n",
        "        missing = [k for k, v in found_files.items() if v is None]\n",
        "        if missing:\n",
        "            print(f\"Error: Missing required files: {missing}\")\n",
        "            print(f\"Uploaded files: {list(uploaded.keys())}\")\n",
        "            return\n",
        "\n",
        "        print(\"\\nFiles found:\")\n",
        "        for file_type, filename in found_files.items():\n",
        "            print(f\"- {file_type}: {filename}\")\n",
        "\n",
        "        # Load JSON contents\n",
        "        json_contents = {}\n",
        "        for file_type, filename in found_files.items():\n",
        "            with open(filename, 'r', encoding='utf-8') as f:\n",
        "                try:\n",
        "                    json_contents[file_type] = f.read()\n",
        "                except Exception as e:\n",
        "                    print(f\"Error reading {file_type} file: {str(e)}\")\n",
        "                    return\n",
        "\n",
        "        print(\"\\nStarting academic assistance workflow...\")\n",
        "        llm = NeMoLLaMa(NEMOTRON_4_340B_INSTRUCT_KEY)\n",
        "        coordinator_output, output = await run_all_system(\n",
        "            json_contents['profile'],\n",
        "            json_contents['calendar'],\n",
        "            json_contents['task']\n",
        "        )\n",
        "        return coordinator_output, output\n",
        "\n",
        "    except Exception as e:\n",
        "        print(f\"\\nError: {str(e)}\")\n",
        "        print(\"\\nDetailed error information:\")\n",
        "        import traceback\n",
        "        print(traceback.format_exc())\n",
        "        return None, None\n",
        "\n",
        "# Run the system\n",
        "coordinator_output, output = await load_json_and_test()"
      ],
      "metadata": {
        "id": "2LM2FEtoFUNa"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## LLM Output\n",
        "\n",
        "code for show output in markdown format"
      ],
      "metadata": {
        "id": "_WM90t6x272I"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Your text content comes as a JSON string\n",
        "try:\n",
        "    # Parse the JSON string\n",
        "    if isinstance(output, str):\n",
        "        json_content = json.loads(output)\n",
        "    else:\n",
        "        json_content = output\n",
        "\n",
        "    # Extract just the plan content and clean it\n",
        "    plan_content = json_content.get('plan', '')\n",
        "\n",
        "    # Remove unnecessary characters and formats\n",
        "    plan_content = plan_content.replace('\\\\n', '\\n')  # Convert \\n string to actual newlines\n",
        "    plan_content = plan_content.replace('\\\\', '')     # Remove remaining backslashes\n",
        "    plan_content = re.sub(r'\\{\"plan\": \"|\"\\}$', '', plan_content)  # Remove JSON wrapper\n",
        "\n",
        "    # Create a console instance\n",
        "    console = Console()\n",
        "\n",
        "    # Create a markdown object with the cleaned content\n",
        "    md = Markdown(plan_content)\n",
        "\n",
        "    # Create a panel with the markdown content\n",
        "    panel = Panel(md, title=\"LLM Output\", border_style=\"blue\")\n",
        "\n",
        "    # Print the formatted content\n",
        "    console.print(panel)\n",
        "\n",
        "except Exception as e:\n",
        "    print(f\"Error formatting output: {e}\")\n",
        "    print(\"Raw output:\", output)"
      ],
      "metadata": {
        "id": "ZJis5pt42wGU"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Example of output (as the system could take a long time to execute)"
      ],
      "metadata": {
        "id": "cFFtxoKIhXWw"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "![image.png]()"
      ],
      "metadata": {
        "id": "aPS2uHKadIv1"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "![image.png]()"
      ],
      "metadata": {
        "id": "pxyyu2dZdWls"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "#🚀 Example of use cases"
      ],
      "metadata": {
        "id": "8D8fjwBrdsmP"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# I) Time Management & Focus Emergency\n",
        "## User Request: \"Help! I have a physics exam tomorrow but can't stop checking TikTok. Need to study 5 chapters tonight"
      ],
      "metadata": {
        "id": "lZfJnerOfQmc"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "System: Let me analyze your situation and create a focused study plan.\n",
        "\n",
        "Coordinator Analysis:\n",
        "- Emergency exam situation\n",
        "- Social media distraction issue\n",
        "- Large content volume\n",
        "- Time constraint\n",
        "\n",
        "Activated Agents:\n",
        "✓ Planner Agent: Create time-blocked schedule\n",
        "✓ NoteWriter Agent: Generate quick study guides\n",
        "✓ Advisor Agent: Provide focus strategies\n",
        "\n",
        "Generated Support Plan:\n",
        "\n",
        "🕒 EMERGENCY SCHEDULE (Next 8 Hours):\n",
        "\"Here's your personalized focus-emergency plan:\n",
        "\n",
        "7:00-7:30 PM: SETUP\n",
        "- Delete TikTok (temporarily!)\n",
        "- Set up study space\n",
        "- Quick energy snack\n",
        "- Download offline PDF chapters\n",
        "\n",
        "7:30-9:30 PM: POWER HOUR 1\n",
        "- Chapters 1-2 speed review\n",
        "- 25/5 min focus/break cycles\n",
        "- Stand up every switch\n",
        "- Quick voice recordings for key concepts\n",
        "\n",
        "9:30-9:45 PM: RESET BREAK\n",
        "- Walk around\n",
        "- Hydrate\n",
        "- No phone zone!\n",
        "\n",
        "[Schedule continues...]\n",
        "\n",
        "🧠 FOCUS HACKS:\n",
        "- Forest app activated\n",
        "- Phone in another room\n",
        "- Study position changes every 30 mins\n",
        "- Voice record summaries instead of writing\n",
        "\n",
        "📚 EMERGENCY STUDY GUIDE:\n",
        "[Condensed chapter summaries...]\n",
        "\n",
        "🚨 ANTI-DISTRACTION PROTOCOLS:\n",
        "If urge to check phone hits:\n",
        "1. Do 10 jumping jacks\n",
        "2. Change study position\n",
        "3. Record voice note instead\n",
        "4. Switch to standing desk mode\n"
      ],
      "metadata": {
        "id": "eP_UlAIfexrc"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# II) Course Material Processing\n",
        "## User Request: \"Turn my 3-hour lecture recording and 45 slides into quick study notes. I'm a visual learner with ADHD.\""
      ],
      "metadata": {
        "id": "C2RNNz00fiRB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "System: Analyzing your learning style and content for optimal processing.\n",
        "\n",
        "Generated Support:\n",
        "\n",
        "🎨 VISUAL STUDY GUIDE:\n",
        "\"Created these ADHD-friendly study materials:\n",
        "\n",
        "MIND MAP SUMMARIES:\n",
        "[Visual concept maps for each major topic]\n",
        "\n",
        "COLOR-CODED QUICK REFERENCES:\n",
        "🟦 Core Concepts\n",
        "🟨 Key Examples\n",
        "🟩 Formula Applications\n",
        "🟥 Common Mistakes\n",
        "\n",
        "MINI VISUAL CHUNKS:\n",
        "- 3-minute concept videos\n",
        "- Animated formula breakdowns\n",
        "- Visual memory hooks\n",
        "- Practice problem flowcharts\n"
      ],
      "metadata": {
        "id": "aQcaqtVcezOE"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### We store course materials in Vector Database and this will be used in our future plan ATLAS Assistant Application"
      ],
      "metadata": {
        "id": "9P-WvHdUgeFS"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "![image.png]()"
      ],
      "metadata": {
        "id": "g5KcWJ-BgYK8"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "# III) Weekly Planning & Goal Setting\n",
        "## User Request: \"Help me plan next week - I have football practice Mon/Wed/Fri, two assignments due Thursday, and need to start preparing for midterms.\""
      ],
      "metadata": {
        "id": "cHSXj4cJfpFJ"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "\n",
        "System: Creating a balanced plan that works with your energy levels and commitments.\n",
        "\n",
        "Generated Plan:\n",
        "\n",
        "⚡ ENERGY-OPTIMIZED SCHEDULE:\n",
        "\n",
        "MONDAY:\n",
        "Pre-Practice (9-11 AM):\n",
        "- Assignment 1 research\n",
        "- Quick progress tracker update\n",
        "Post-Practice (7-9 PM):\n",
        "- Light review sessions\n",
        "- Assignment outline\n",
        "\n",
        "TUESDAY (Peak Focus Day):\n",
        "Morning Power Block (9-12):\n",
        "- Assignment 1 completion\n",
        "- Midterm topic list\n",
        "Afternoon (2-5):\n",
        "- Assignment 2 deep work\n",
        "- Create study guides\n",
        "\n",
        "[Schedule continues with similar detail...]\n",
        "\n",
        "🎯 STRATEGY NOTES:\n",
        "\"I've noticed you perform better with morning study sessions after practice days. I've scheduled intensive work during these peak energy times.\n",
        "\n",
        "Each day includes:\n",
        "- Energy level tracking\n",
        "- Flexibility blocks\n",
        "- Recovery time\n",
        "- Progress check-ins"
      ],
      "metadata": {
        "id": "SBapV_E9fuYB"
      }
    },
    {
      "cell_type": "markdown",
      "source": [],
      "metadata": {
        "id": "qn_xtOfTf58p"
      }
    }
  ]
}
